Tensorflow docker 图像不使用 GPU 进行推理/预测

Question

我的要求：

使用 tensorflow 在 GPU 上运行推理任务以进行对象检测。

当前状态：

我正在使用 AWS GPU 实例 (p2.xlarge) 进行训练和推理。 训练部分在 GPU 上运行良好。这里没问题。 （显卡：Tesla M60）

为了获得预测，我创建了一个烧瓶服务器，它封装了 tensorflow 检测，并为其添加了一些额外的逻辑。我打算将此服务（Flask + tensorflow）部署为 docker 容器。我使用的基础镜像是tensorflow/tensorflow:1.12.0-gpu-py3。我的dockerfile 看起来像这样：

FROM tensorflow/tensorflow:1.12.0-gpu-py3
COPY ./app /app
COPY ./requirements.txt /app
RUN pip3 install -r /app/requirements.txt
RUN mkdir /app/venv
WORKDIR /app
RUN export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
ENTRYPOINT ["python3", "/app/main.py"]
ENV LISTEN_PORT 8080
EXPOSE 8080

我可以通过以下方式部署它：

docker run --runtime=nvidia --gpus all --name <my-long-img-name> 
-v <somepath>:<anotherpath> -p 8080:8080 -d <my-long-img-name>

并成功从邮递员调用 8080 端口上的端点。

基本上，我的意思是所有驱动程序都已正确设置。

flask 中的一个端点是这样的：（用于测试是否使用 GPU）

@app.route("/testgpu", methods=["GET"])
def testgpu():
    import tensorflow as tf
    with tf.device('/gpu:0'):
        a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
        b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
        c = tf.matmul(a, b)

    with tf.Session() as sess:
        print (sess.run(c))

当我调用这个端点时，我没有收到任何错误（如果没有检测到 gpu，它会抛出错误）。这意味着为这个 sn-p 检测到 gpu。耶！！

我还将这两行添加到我的主要代码执行流程中：

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

它输出：

Local devices : 
[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 17661279486087266140
, name: "/device:XLA_GPU:0"
device_type: "XLA_GPU"
memory_limit: 17179869184
locality {
}
incarnation: 9205152708262911170
physical_device_desc: "device: XLA_GPU device"
, name: "/device:XLA_CPU:0"
device_type: "XLA_CPU"
memory_limit: 17179869184
locality {
}
incarnation: 3134142118233627849
physical_device_desc: "device: XLA_CPU device"
, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 7447009690
locality {
  bus_id: 1
  links {
  }
}
incarnation: 6613138223738633761
physical_device_desc: "device: 0, name: Tesla M60, pci bus id: 0000:00:1e.0, compute capability: 5.2"
]

又是一次，GPU 被检测到了。

甚至来自 tensorflow 的日志也占用了 GPU。

2019-11-18 08:45:29.944580: I tensorflow/core/common_runtime/gpu/gpu_device.cc:982] Device interconnect StreamExecutor with strength 1 edge matrix:
2019-11-18 08:45:29.944603: I tensorflow/core/common_runtime/gpu/gpu_device.cc:988]      0 
2019-11-18 08:45:29.944611: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1001] 0:   N 
2019-11-18 08:45:29.944721: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 7101 MB memory) -> physical GPU (device: 0, name: Tesla M60, pci bus id: 0000:00:1e.0, compute capability: 5.2)

这里一切似乎都很顺利，但 GPU 应该运行的主要部分并没有接受它。它正在使用CPU。还有另一个端点（比如/getpredictions）以及上面提到的/testgpu，它运行检测并返回输出。

问题： 每当我在端口 8080 上从邮递员调用 /getpredictions 而不是使用 GPU 时，它都会占用 CPU 并在大约 30 多秒内返回输出。

这里有什么遗漏吗？有什么解决方法吗？

如果我需要在问题中添加更多信息，请告诉我。

Answer 1

从文档中你应该在运行容器时添加 gpu 选项 像这样：

docker run -it --gpus all -p 8888:8888 tensorflow/tensorflow:latest-gpu-jupyter