Fix batch-size on resume for multi-gpu

[NanoCode012]

Fixes #1936

Tested on my own and with author of the mentioned Issue.

Command:

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

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Improvement in resume training functionality for YOLOv5.

📊 Key Changes

• Changed the train.py script to update opt.batch_size properly when resuming training.

🎯 Purpose & Impact

• The purpose of this change is to ensure that the batch size used when resuming training from a checkpoint is set correctly to the total batch size (opt.total_batch_size).
• This will prevent potential inconsistencies with batch sizes when training is paused and resumed, leading to more reliable and stable training runs for users. 🛠️

This update is particularly beneficial for users who conduct long training sessions that may need to be paused and resumed due to various reasons, such as hardware limitations or scheduling constraints. 🔄

[glenn-jocher]

@NanoCode012 I was looking back at the v4.0 release commit 69be8e7 and it seems I might have caused this by moving one line here in train.py:

[glenn-jocher]

I had a strange situation where I was running some v3.1 trainings multi-GPU, and then when I resumed one I got CUDA OOM (my original run was very close to OOM already, and perhaps resume used just a little more memory), so I modified a few of these lines to allow for a --batch command on resume. I think I forgot to reset this one back to default.

[glenn-jocher]

@NanoCode012 do you think updating the PR to simply revert the earlier change would fix this, or do you think the current PR is best?

[NanoCode012]

I don't think reverting fixes the problem. The problem is on-resume. The batch-size should be the previous "total-batchsize" for this

yolov5/train.py

Line 491 in b75c432

opt.total_batch_size = opt.batch_size

before it would be divided here

yolov5/train.py

Line 499 in b75c432

opt.batch_size = opt.total_batch_size // opt.world_size

Edit: If you move that line up, I think you still need this PR.

[glenn-jocher]

@NanoCode012 ok understood! Merging PR.