如何在多个 GPU 上使用 tensorflow 并行推理？

Question

我在部署我的 tensorflow 模型时使用 Tensorflow-serving。我在可用的服务器上有多个 GPU，但是，截至目前在推理期间，只使用了一个 GPU。

我现在的想法是，为了并行化大量图像的分类，为每个可用的 GPU 生成一个张量流服务图像，并并行 "workers" 从生成器抓取图像，发出请求并等待答案。然后从生成器中获取新图像，依此类推。这意味着我必须实现自己的数据处理程序，但这似乎是可以实现的。

我在 TensorFlow Serving Batching 中阅读了一些关于 SharedBatchScheduler 的内容，但我不知道这是否有用，请查看更多内容。

总的来说，我对 tensorflow-serving 还很陌生，我想知道这是否是实现我想要的最直接的方法。

提前感谢任何help/suggestions！

编辑： 感谢澄清问题：我知道 311 问题，github.com/tensorflow/serving/issues/311。有人有解决此问题的方法吗？

Answer 1

用docker和nvidia-docker 2.0完全可行（从docker run --runtime=nvidia ...和issue来看，他们用的是第一个版本）。我确实尝试过使用多个 GPU 和服务；但是，并没有在多个 GPU 上结束运行。

尽管如此，我有一台带 4 个 GPU 的主机，目前为每个具有 Tensorflow 运行的自定义图像调度 1 个 GPU 进行训练，以便每个用户都可以在隔离环境中使用 GPU。以前我使用 Kubernetes 进行设备配置和容器管理，但这对我需要做的事情来说太过分了。目前，我正在使用 docker-compose 来完成所有魔术。这是一个例子：

version: '3'
services:
    lab:
        build: ./tensorlab
        image: centroida/tensorlab:v1.1
        ports:
            - "30166:8888"
            - "30167:6006"
        environment:
            NVIDIA_VISIBLE_DEVICES: 0,1,2
       ...

这里的关键部分是NVIDIA_VISIBLE_DEVICES变量，其中一个GPU的索引对应nvidia-smi

的输出

如何在多个 GPU 上使用 tensorflow 并行推理？

How to parallelize inference with tensorflow serving on multiple GPUs?

python

tensorflow

tensorflow-serving