Add tensor hooks and 10.0 gradient clipping

[UnglvKitDe]

Improves the stability issues as described in #8578

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Improving training stability with gradient regularization and NaN mitigation.

📊 Key Changes

• 🧊 Added a hook to convert NaN values to zero during training, reinforcing model stability.
• ✂️ Introduced gradient clipping to prevent excessively large gradients, which can corrupt model weights.

🎯 Purpose & Impact

• 🛡️ The NaN conversion to zero avoids potential crashes or instability in training by handling undefined numerical values gracefully.
• 💪 Gradient clipping safeguards the training process by keeping gradients within a manageable range, promoting healthier and more stable weight updates.
• ⚙️ These enhancements are geared towards making the YOLOv5 training process more robust and reliable, beneficial for users looking to train models without encountering numerical issues.

[UnglvKitDe]

@glenn-jocher I have to change the part with the hook, that is not so right I noticed. Srry.

[glenn-jocher]

@glenn-jocher I have to change the part with the hook, that is not so right I noticed. Srry.

Is this PR ready to test or do you still need to make changes?

[UnglvKitDe]

@glenn-jocher Yes now, I removed the unnecessary line of code. It would not have changed, but it is not necessary. Thanks a lot!

[UnglvKitDe]

@glenn-jocher A few insights: I've been running ~50 trainings on my private dataset with this new branch the last few days. So far it looks like this version is much more stable & the rSeed plays a less important role. No unstable trainings as described above have occurred so far. Before sometimes 3 out of 9 trainings were unstable. I know that is just my dataset and not a general test :)

[glenn-jocher]

@UnglvKitDe I'd like to close this out but we need a few changes. If this improves default trainings with little cost/resources penalty then we should enable it by default. We want to avoid adding argparser arguments if at all possible, we already have too many. Can you profile training with this branch and with master to compare mAP and training time and CUDA memory utilization?

[glenn-jocher]

@UnglvKitDe also why do we have to unscale optimizer first? This seems like a costly operation, wouldn't it make more sense to scale the gradient clipping value by the scaler settings?

[UnglvKitDe]

@UnglvKitDe I'd like to close this out but we need a few changes. If this improves default trainings with little cost/resources penalty then we should enable it by default. We want to avoid adding argparser arguments if at all possible, we already have too many. Can you profile training with this branch and with master to compare mAP and training time and CUDA memory utilization?

@glenn-jocher At the moment I unfortunately don't have the resource to run a full coco training. My own dataset has only a few hundred images, coco ~120k. Srry! Can you test this? If I can help you in any other way, please let me know :)

[UnglvKitDe]

@UnglvKitDe also why do we have to unscale optimizer first? This seems like a costly operation, wouldn't it make more sense to scale the gradient clipping value by the scaler settings?

@glenn-jocher As described here, it is necessary to unscale them because otherwise you will get the wrong gradients. II did not notice any significant longer training on my data set. But (as said above) I have only a small dataset. Maybe it also improves the training and leads to better results (by a more stable training).

[glenn-jocher]

@UnglvKitDe thanks, I'll review the link and run some tests this weekend.

[glenn-jocher]

@UnglvKitDe it looks like the PR does two separate things, first replace NaN with 0.0 and then separately clip gradients. Do you know which change resulted into the improvements you saw? Are you able to determine which was more important?

[glenn-jocher]

@UnglvKitDe I'm testing in Colab now. I see same time but slightly worse mAP for PR. I wonder if the clipping is too tight. We have 10.0 now, maybe put 30.0 or 100.0?

!git clone https://github.com/UnglvKitDe/yolov5-1 -b fix/grad_inf_nan  # clone
%cd yolov5-1
%pip install -qr requirements.txt  # install
!python train.py --img 640 --batch 16 --epochs 30 --data coco128.yaml --weights yolov5s.pt --cache

%cd ..
!git clone https://github.com/ultralytics/yolov5
%cd yolov5
!python train.py --img 640 --batch 16 --epochs 30 --data coco128.yaml --weights yolov5s.pt --cache

[glenn-jocher]

@UnglvKitDe confirm raising the clip max to 100.0 solves the issue of lower mAP. You think this might be too high?

[UnglvKitDe]

@UnglvKitDe it looks like the PR does two separate things, first replace NaN with 0.0 and then separately clip gradients. Do you know which change resulted into the improvements you saw? Are you able to determine which was more important?
@glenn-jocher For me, both in combination have improved the ressults.

[UnglvKitDe]

@glenn-jocher About your changes: I don't know to what extent the parameters may be None. Actually, this makes no sense. But in a forum this was advised once, but it can also be an old & solved bug. I have seen values for max_value from 2.0 to 5.0. 10 was already chosen very high by me. I find 100, is already very very high...

[UnglvKitDe]

@glenn-jocher I think that depends on the context, which height is needed, so I wanted to keep it configurable. I agree with you though that we have too many argparser arguments. Maybe we put this in a config file (hyp*.yaml or model *.yaml ?).

[glenn-jocher]

@UnglvKitDe reduced gradient clipping to 10.0 after some experiments found not much difference with master.

PR is merged. Thank you for your contributions to YOLOv5 🚀 and Vision AI ⭐

[UnglvKitDe]

@glenn-jocher Thx :)

[glenn-jocher]

@UnglvKitDe I observed erratic training behavior (green line) with the nan_to_num hook in classifier branch (I added it there also), so I'm going to remove it from master.

[UnglvKitDe]

@glenn-jocher Mh interesting, I have never seen anything like this. You have taken my code except for the None check I see. I do not know if this is the reason. But actually none of the weights should be None. And if they are, an exception should be thrown...

[glenn-jocher]

@UnglvKitDe I agree, the None checks should not be necessary, and the issue should be investigated further to ensure all weights are properly initialized and exceptions are handled appropriately. Thank you for bringing this to my attention.