在《TensorFlow之目标检测API接口调试(超详细)》文章中,介绍了tensorflow API接口调用方法,详细描述了训练的过程,训练的时候用了单GPU训练,tensorflow是支持多GPU训练的,因此,介绍一下多GPU训练的方法。

 

很简单,只需要修改两个地方:

第一个参数是train.py文件,找到以下参数

 

flags.DEFINE_integer('num_clones', 1, 'Number of clones to deploy per worker.')
默认为1,如果自己有多个GPU,可以将参数修改为GPU数量,比如说2、4等;当然也可以在运行train.py的时候以参数的形式传入。

 

第二个参数是训练配置文件,如rfcn_resnet101_coco.config,根据自己使用的具体网络不同,在train_config中,找到batch_size参数,默认值为1,可修改为大点的参数,理论上

大点比较好,并且要求是GPU数量的整数倍,比如设置为8、16、32等,这个参数根据显存来决定,显存越大,可设置的越大。

 

在batch_size一定的情况下,GPU数量越多,训练速度越快

 

修改完上述两个参数后,就可以调用多个GPU来训练了。

 

注意batch_size的设置要合理,如果设置过大,会报如下所示的错误

ResourceExhaustedError (see above for traceback): OOM when allocating tensor with shape[8,2048,36,64] and type float on /job:localhost/replica:0/task:0/device:GPU:1 by allocator GPU_1_bfc

如果说想要指定GPU号,则在train.py中添加以下代码

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'

表示调用第1块和第4块显卡训练