最近使用多卡训练YOLOV7,一开始使用的是DP,可以正常训练,但是看到使用DDP可以加速训练,果断使用,但是总报错,之前YOLOV5使用没问题,所以就照着YOLOv5改了改,断断续续改了俩天终于成功了!!!
话不多说,直接上代码:
train.py 最前面加上
LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html
RANK = int(os.getenv('RANK', -1))
WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
在main函数有如下改动
# Set DDP variables
# opt.world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
# opt.global_rank = int(os.environ['RANK']) if 'RANK' in os.environ else -1
opt.global_rank = RANK
opt.local_rank = LOCAL_RANK
opt.world_size = WORLD_SIZE
我觉得主要原因在于opt.local_rank=int(os.getenv(‘RANK’, -1))这一行!!!