tensorflow - 如何使用 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"])