Repository: ultralytics/yolov5

Issue:
CUDA device index error in distributed training
## Question

When running the distributed training, the following line is reporting a CUDA device index error

https://github.com/ultralytics/yolov5/blob/276b6745132f2da0641706124c466b69f21b670c/utils/torch_utils.py#L37

Sample script when launching the training

```
# On master machine 0
$ python -m torch.distributed.launch --nproc_per_node G --nnodes N --node_rank 0 --master_addr "192.168.1.1" --master_port 1234 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights sample_weight.pt  --device 0,...G-1''
```

```
# On machine R
$ python -m torch.distributed.launch --nproc_per_node G --nnodes N --node_rank R --master_addr "192.168.1.1" --master_port 1234 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights sample_weight.pt  --device 0,...G-1''
```

All instances have the sample number of GPU per machine.

When running the script above,

- master instance is running with no error
- machine R is complaining about the CUDA device index error

Proposed fix: changing the following lines from `RANK` to `LOCAL_RANK`

[1] https://github.com/ultralytics/yolov5/blob/276b6745132f2da0641706124c466b69f21b670c/train.py#L102
[2] https://github.com/ultralytics/yolov5/blob/276b6745132f2da0641706124c466b69f21b670c/train.py#L114
[3] https://github.com/ultralytics/yolov5/blob/276b6745132f2da0641706124c466b69f21b670c/train.py#L211

Context info:

- all instances running in docker container, with
    - NVIDIA Release 21.03 (build 21060478)
    - PyTorch Version 1.9.0a0+df837d0

Question: 

- does the proposed fix above look good?



Discussion:

Resolving patch:
diff --git a/train.py b/train.py
--- a/train.py
+++ b/train.py
@@ -99,7 +99,7 @@ def train(hyp,  # path/to/hyp.yaml or hyp dictionary
     plots = not evolve  # create plots
     cuda = device.type != 'cpu'
     init_seeds(1 + RANK)
-    with torch_distributed_zero_first(RANK):
+    with torch_distributed_zero_first(LOCAL_RANK):
         data_dict = data_dict or check_dataset(data)  # check if None
     train_path, val_path = data_dict['train'], data_dict['val']
     nc = 1 if single_cls else int(data_dict['nc'])  # number of classes
@@ -111,7 +111,7 @@ def train(hyp,  # path/to/hyp.yaml or hyp dictionary
     check_suffix(weights, '.pt')  # check weights
     pretrained = weights.endswith('.pt')
     if pretrained:
-        with torch_distributed_zero_first(RANK):
+        with torch_distributed_zero_first(LOCAL_RANK):
             weights = attempt_download(weights)  # download if not found locally
         ckpt = torch.load(weights, map_location=device)  # load checkpoint
         model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
@@ -208,7 +208,7 @@ def train(hyp,  # path/to/hyp.yaml or hyp dictionary
 
     # Trainloader
     train_loader, dataset = create_dataloader(train_path, imgsz, batch_size // WORLD_SIZE, gs, single_cls,
-                                              hyp=hyp, augment=True, cache=opt.cache, rect=opt.rect, rank=RANK,
+                                              hyp=hyp, augment=True, cache=opt.cache, rect=opt.rect, rank=LOCAL_RANK,
                                               workers=workers, image_weights=opt.image_weights, quad=opt.quad,
                                               prefix=colorstr('train: '))
     mlc = int(np.concatenate(dataset.labels, 0)[:, 0].max())  # max label class


Test patch: