如何使用 Tensorflow 对象检测 API 启用多 GPU 训练

Question

我正在尝试使用 TensorFlow 对象检测 API 执行多 GPU 训练。

我在 NVIDIA-SMI 中看到的是，实际上只使用了 1 个 GPU。提供的其他 3 个 GPU 已加载 GPU 进程，但内存使用量为 300MB，利用率始终为 0%

我正在使用在 COCO 上预训练的基于 SSD MobileNetV1 的网络，然后使用我的自定义数据集对其进行训练。

我预计当我为 Tensorflow 提供更多 GPU 时，该框架实际上会使用它们来加速训练。

Answer 1

对于 Tensorflow 2.2.0 对象检测 API，当您运行 model_main_tf2.py 时，启用此标志：

python model_main_tf2.py --num_workers=2

对于 --num_workers > 1 的任何整数，tensorflow 使用所有可用的 GPU，如果您只想使用部分 GPU，则必须编辑此 model_main_tf2.py 文件，其中指定策略，同时将 num_workers 保留为默认值 1。例如，使用机器的第一个和第二个 GPU：

strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])

Answer 2

不幸的是，目前还不可能。您可以按照此处所述运行多 GPU 推理：https://github.com/tensorflow/models/issues/6611#issuecomment-508507576。但出于训练目的，这是不可能的。我听说开发人员正在致力于将 bas 迁移到 keras。这可以实现多 GPU 训练。