在 aws sagemaker 中配置 GPU，使用 keras 和 tensorflow 作为后端

Question

我是 aws sagemaker 的新手。我正在尝试使用具有 GPU 支持的 keras 在 aws sagemaker 中设置模型。下面给出了用于推断模型的 docker 基础镜像

FROM tensorflow/tensorflow:1.10.0-gpu-py3

RUN apt-get update && apt-get install -y --no-install-recommends nginx curl
...

这是我用来检查 Flask 中的 keras 是否识别 GPU 的 keras 代码。

import keras
@app.route('/ping', methods=['GET'])
def ping():

    keras.backend.tensorflow_backend._get_available_gpus()

    return flask.Response(response='\n', status=200,mimetype='application/json')

当我使用 GPU 在 sagemaker 中启动笔记本实例时，keras 代码会显示可用的 GPU。那么，为了在推理阶段（模型）访问 GPU，除了张量流 GPU 基础映像之外，我是否还需要在 docker 文件中安装任何其他库？

提前致谢。

Answer 1

您不需要安装任何其他东西。Keras 依赖 TensorFlow 进行 GPU 检测和配置。

唯一值得注意的是训练时如何使用多个GPU。我建议将“gpu_count”作为超级参数传递，并进行如下设置：

from keras.utils import multi_gpu_model
model = Sequential()
model.add(...)
...
if gpu_count > 1:
    model = multi_gpu_model(model, gpus=gpu_count)
model.compile(...)