Repository: ultralytics/yolov5

Issue:
wrong batch size after --resume on multiple GPUs
## 🐛 Bug
After running  a training session on multiple GPUs, batch_size is read wrongly from opt.yaml causing an error.

## To Reproduce (REQUIRED)
Run this line:
```
python -m torch.distributed.launch --nproc_per_node 2 train.py --batch-size 30 --data coco128.yaml --weights yolov5s.pt --device 1,2
```
* Notice that we use batch-size=30 --> 15 for each GPU.

Press Ctrl+C to stop the session.
Now, try to resume:
```
python -m torch.distributed.launch --nproc_per_node 2 train.py --batch-size 30 --data coco128.yaml --weights yolov5s.pt --device 1,2 --resume
```

Output:
```
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
*****************************************
github: Traceback (most recent call last):
  File "train.py", line 492, in <module>
    device = select_device(opt.device, batch_size=opt.batch_size)
  File "/user/detection/tests/yolov5/utils/torch_utils.py", line 68, in select_device
    assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
AssertionError: batch-size 15 not multiple of GPU count 2
up to date with https://github.com/ultralytics/yolov5 ✅
Resuming training from ./runs/train/exp/weights/last.pt
Traceback (most recent call last):
  File "train.py", line 492, in <module>
    device = select_device(opt.device, batch_size=opt.batch_size)
  File "/user/detection/tests/yolov5/utils/torch_utils.py", line 68, in select_device
    assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
AssertionError: batch-size 15 not multiple of GPU count 2
Traceback (most recent call last):
  File "/usr/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/user/venv/yolov5/lib/python3.8/site-packages/torch/distributed/launch.py", line 260, in <module>
    main()
  File "/user/venv/yolov5/lib/python3.8/site-packages/torch/distributed/launch.py", line 255, in main
    raise subprocess.CalledProcessError(returncode=process.returncode,
subprocess.CalledProcessError: Command '['/user/venv/yolov5/bin/python', '-u', 'train.py', '--local_rank=1', '--batch-size', '30', '--data', 'coco128.yaml', '--weights', 'yolov5s.pt', '--device', '1,2', '--resume']' returned non-zero exit status 1.

```

## Expected behavior
The training should resume correctly, with the right batch size.

## Environment

 - OS: [Ubuntu 16.04]
 - GPU [GTX 1080 Ti (x2)]


## Additional context
--


Discussion:
I'm thinking that we need to modify the below a bit . 
https://github.com/ultralytics/yolov5/blob/b75c432ea0b35b44ba364bc6acadfc85f7dddf73/train.py#L480

to

```python
opt.cfg, opt.weights, opt.resume, opt.batch_size, opt.global_rank, opt.local_rank  = '', ckpt, True, opt.total_batch_size, *apriori
```

@AmirSa7 , could you try to modify it and tell me how it goes?

Edit: I tested it quickly and it works.

Commands:

```bash
python -m torch.distributed.launch --master_port 9963 --nproc_per_node 2 train.py --data coco128.yaml --cfg yolov5s.yaml --weights yolov5s.pt --epochs 3 --batch-size 64 --device 3,4
python -m torch.distributed.launch --master_port 9963 --nproc_per_node 2 train.py --resume
```
> I'm thinking that we need to modify the below a bit .
> https://github.com/ultralytics/yolov5/blob/b75c432ea0b35b44ba364bc6acadfc85f7dddf73/train.py#L480
> 
> to
> 
> ```python
> opt.cfg, opt.weights, opt.resume, opt.batch_size, opt.global_rank, opt.local_rank  = '', ckpt, True, opt.total_batch_size, *apriori
> ```
> 
> @AmirSa7 , could you try to modify it and tell me how it goes?
> 
> Edit: I tested it quickly and it works.
> 
> Commands:
> 
> ```shell
> python -m torch.distributed.launch --master_port 9963 --nproc_per_node 2 train.py --data coco128.yaml --cfg yolov5s.yaml --weights yolov5s.pt --epochs 3 --batch-size 64 --device 3,4
> python -m torch.distributed.launch --master_port 9963 --nproc_per_node 2 train.py --resume
> ```

Yup, it works :+1:  Thanks.

Resolving patch:
diff --git a/train.py b/train.py
--- a/train.py
+++ b/train.py
@@ -477,7 +477,7 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
         apriori = opt.global_rank, opt.local_rank
         with open(Path(ckpt).parent.parent / 'opt.yaml') as f:
             opt = argparse.Namespace(**yaml.load(f, Loader=yaml.FullLoader))  # replace
-        opt.cfg, opt.weights, opt.resume, opt.global_rank, opt.local_rank = '', ckpt, True, *apriori  # reinstate
+        opt.cfg, opt.weights, opt.resume, opt.batch_size, opt.global_rank, opt.local_rank = '', ckpt, True, opt.total_batch_size, *apriori  # reinstate
         logger.info('Resuming training from %s' % ckpt)
     else:
         # opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml')


Test patch: