【Tensorflow object detection API】多GPU卡并行训练+Ubuntu 16.04(完美运行,教程步骤清晰!!)

  • 问题描述:我的电脑有3块2080ti的GPU,在采用Tensorflow object_detection API原有的model_main.py进行训练时,源于tensorflow本身性质,直接占满3块GPU,本来以为这样就是利用了多卡进行训练(结果还是想得太简单);
  • 出错:可是将batch_size设置成128就会显存溢出,照理GPU的性能是够用的
  • 解决思路
    利用nvidia-smi命令观察GPU使用情况,发现虽然是占满了3张卡,可是默认只使用了1张卡在训练,并没有利用3卡训练,如下图:
    在这里插入图片描述
解决办法:
  1. 不能采用model_main.py训练,因为Estimator分发策略不适用于tf.contrib.slim。所以没有关于分布式训练,多卡训练的部分。
  2. 具体步骤有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
  1. 采用如上步骤后,可以看到完美的利用了3张GPU,可以将batch_size调大了!!
    在这里插入图片描述

猜你喜欢

转载自blog.csdn.net/qq_43348528/article/details/107065506