Mixed Precision Training

[JSabadin]

Mixed Precision Training Implementation

This pull request introduces mixed precision training to improve the performance and efficiency of our model training.

Key Changes

• Mixed precision training has been implemented, leveraging the automatic mixed precision (AMP) feature available in PyTorch.
• Testing results:
  • Prior to this implementation, the detection model took:
    • 10.5 minutes per epoch on a batch size of 16 (bs16)
    • 10.5 minutes per epoch on a batch size of 32 (bs32) due to high GPU utilization and full VRAM.
  • After implementing mixed precision training, the detection epoch time reduced to:
    • 8.5 minutes per epoch on a batch size of 32 (bs32), providing a significant performance boost.

Limitations

• This implementation has only been tested on detection models. It is recommended to perform further testing on other models.
• Some loss functions do not support mixed precision training. In these cases, the following should be applied:
  • Use torch.amp.autocast(..., disable=True) to disable AMP for specific loss functions.

Benefits of Mixed Precision Training

• Reduced training time: By utilizing half-precision for certain operations, training time has been significantly reduced while maintaining model accuracy.
• Better hardware utilization: Improved GPU memory usage, allowing for larger batch sizes without running into VRAM issues.

[kozlov721]

LGTM