Ankle/yolov5_ZWgra
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

Update DDP backend if dist.is_nccl_available() (#3705)

Browse Source
This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-06-20 17:15:42 +02:00
1 parent fbf41e0913
commit e8810a53e8
1 file changed
+1 -1
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+1 -1
View File
@@ -539,7 +539,7 @@ def main(opt):
assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command' assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
torch.cuda.set_device(LOCAL_RANK) torch.cuda.set_device(LOCAL_RANK)
device = torch.device('cuda', LOCAL_RANK) device = torch.device('cuda', LOCAL_RANK)
dist.init_process_group(backend="gloo", timeout=timedelta(seconds=60)) dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo", timeout=timedelta(seconds=60))
assert opt.batch_size % WORLD_SIZE == 0, '--batch-size must be multiple of CUDA device count' assert opt.batch_size % WORLD_SIZE == 0, '--batch-size must be multiple of CUDA device count'
assert not opt.image_weights, '--image-weights argument is not compatible with DDP training' assert not opt.image_weights, '--image-weights argument is not compatible with DDP training'