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

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
...

这是我用来检查 GPU 是否被 flask 中的 keras 识别的 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，我是否需要在 docker 文件中安装除 tensorflow GPU 基本图像之外的任何其他库？

提前致谢。

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(...)

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

Configuring GPU in aws sagemaker with keras and tensorflow as backend

keras

tensorflow

amazon-sagemaker