🚀 [Add] torch auto mixed precision

yolo/tools/trainer.py

@@ -1,8 +1,10 @@
 import torch
 from loguru import logger
+from torch import Tensor
+from torch.cuda.amp import GradScaler, autocast
 from tqdm import tqdm
 
-from yolo.config.config import TrainConfig
+from yolo.config.config import Config, TrainConfig
 from yolo.model.yolo import YOLO
 from yolo.tools.model_helper import EMA, get_optimizer, get_scheduler
 from yolo.utils.loss import get_loss_function
@@ -22,29 +24,40 @@ class Trainer:
             self.ema = EMA(model, decay=train_cfg.ema.decay)
         else:
             self.ema = None
+        self.scaler = GradScaler()
 
-    def train_one_batch(self, data, targets):
+    def train_one_batch(self, data: Tensor, targets: Tensor, progress: tqdm):
         data, targets = data.to(self.device), targets.to(self.device)
         self.optimizer.zero_grad()
-        outputs = self.model(data)
-        loss = self.loss_fn(outputs, targets)
-        loss.backward()
-        self.optimizer.step()
+
+        with autocast():
+            outputs = self.model(data)
+            loss, loss_item = self.loss_fn(outputs, targets)
+            loss_iou, loss_dfl, loss_cls = loss_item
+
+        progress.set_description(f"Loss IoU: {loss_iou:.5f}, DFL: {loss_dfl:.5f}, CLS: {loss_cls:.5f}")
+
+        self.scaler.scale(loss).backward()
+        self.scaler.step(self.optimizer)
+        self.scaler.update()
+
         if self.ema:
             self.ema.update()
+
         return loss.item()
 
     def train_one_epoch(self, dataloader):
         self.model.train()
         total_loss = 0
-        for data, targets in tqdm(dataloader, desc="Training"):
-            loss = self.train_one_batch(data, targets)
-            total_loss += loss
-        if self.scheduler:
-            self.scheduler.step()
+        with tqdm(dataloader, desc="Training") as progress:
+            for data, targets in progress:
+                loss = self.train_one_batch(data, targets, progress)
+                total_loss += loss
+            if self.scheduler:
+                self.scheduler.step()
         return total_loss / len(dataloader)
 
-    def save_checkpoint(self, epoch, filename="checkpoint.pt"):
+    def save_checkpoint(self, epoch: int, filename="checkpoint.pt"):
         checkpoint = {
             "epoch": epoch,
             "model_state_dict": self.model.state_dict(),