在pytorch分布式训练时,需要指定device_ids
一般我们指定的方式为device_ids=range(cfg.gpus), gpus为我们使用的gpu数量
model = MMDistributedDataParallel(model.cuda(),
device_ids=[cfg.gpus])
但是会出现:RuntimeError: all tensors must be on devices[0]
有可能原因是CUDA_VISIBLE_DEVICES中gpu id和DataParallel中device_ids不同导致的。
解决:
为了每次都能做到匹配,使用动态确定gpu序号的方法。
model = MMDistributedDataParallel(model.cuda(),
device_ids=[torch.cuda.current_device()])
参考:https://blog.csdn.net/qq_16234613/article/details/80948032