- 问题描述:我的电脑有3块2080ti的GPU,在采用Tensorflow object_detection API原有的model_main.py进行训练时,源于tensorflow本身性质,直接占满3块GPU,本来以为这样就是利用了多卡进行训练(结果还是想得太简单);
- 出错:可是将batch_size设置成128就会显存溢出,照理GPU的性能是够用的
- 解决思路:
利用nvidia-smi
命令观察GPU使用情况,发现虽然是占满了3张卡,可是默认只使用了1张卡在训练,并没有利用3卡训练,如下图:
解决办法:
- 不能采用
model_main.py
训练,因为Estimator分发策略不适用于tf.contrib.slim。所以没有关于分布式训练,多卡训练的部分。 - 具体步骤有2步:
(1)采用/legacy/train.py
用2张卡进行训练:
其中CUDA_VISIBLE_DEVICES=0,1
代表使用0号GPU和1号GPU;
CUDA_VISIBLE_DEVICES=0,1 python3 object_detection/legacy/train.py \
--logtostderr \
--pipeline_config_path=object_detection/ssd_model/ssd_mobilenet_v1_pets.config \
--train_dir=object_detection/ssd_model/training6 \
--num_clones=2 --ps_tasks=1
(2)采用legacy/eval.py
用剩下1张卡进行评估:
CUDA_VISIBLE_DEVICES=2 python object_detection/legacy/eval.py \
--logtostderr \
--pipeline_config_path=object_detection/ssd_model/ssd_mobilenet_v1_pets.config \
--checkpoint_dir=object_detection/ssd_model/training6 \
--eval_dir=object_detection/ssd_model/eval6
- 采用如上步骤后,可以看到完美的利用了3张GPU,可以将batch_size调大了!!