How to enable multi GPU training with Tensorflow Object Detection API

user_7846594 0
Viewed 2595

I am attempting to perform Multi-GPU training with the TensorFlow Object Detection API.

What I see in my NVIDIA-SMI is that only 1 GPU is actually being utilized. The other 3 GPUs that are provided have the GPU process loaded to them, but memory usage is at 300MB and utilization sits at 0% at all times

I am using the SSD MobileNetV1 based network pretrained on COCO and then training it with my custom dataset.

I expect that when I provide Tensorflow with more GPUs, the framework will actually use them to speed up training.

tensorflow deep-learning
2 Answers

For Tensorflow 2.2.0 Object Detection API, when you are running model_main_tf2.py, enable this flags:

python model_main_tf2.py --num_workers=2

for any integer for --num_workers > 1, tensorflow uses all available gpus, if you want to use only some of the gpus, you have to edit this model_main_tf2.py file where it specifies the strategy while keeping the num_workers in default 1. This for example, uses first and second gpu of the machine:

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