🔀 [Merge] branch 'TRAIN' into TEST

examples/example_train.py

@@ -28,7 +28,7 @@ def main(cfg: Config):
     # TODO: get_device or rank, for DDP mode
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    trainer = Trainer(model, cfg.hyper.train, device)
+    trainer = Trainer(model, cfg, device)
     trainer.train(dataloader, 10)
 
 

yolo/config/hyper/default.yaml

@@ -1,5 +1,5 @@
 data:
-  batch_size: 4 
+  batch_size: 8 
   shuffle: True
   num_workers: 4
   pin_memory: True

yolo/tools/trainer.py

@@ -1,48 +1,63 @@
 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
 
 
 class Trainer:
-    def __init__(self, model: YOLO, train_cfg: TrainConfig, device):
+    def __init__(self, model: YOLO, cfg: Config, device):
+        train_cfg: TrainConfig = cfg.hyper.train
+
         self.model = model.to(device)
         self.device = device
         self.optimizer = get_optimizer(model.parameters(), train_cfg.optimizer)
         self.scheduler = get_scheduler(self.optimizer, train_cfg.scheduler)
-        self.loss_fn = get_loss_function()
+        self.loss_fn = get_loss_function(cfg)
 
         if train_cfg.ema.get("enabled", False):
             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(),

yolo/utils/loss.py

@@ -17,10 +17,6 @@ from yolo.tools.bbox_helper import (
 )
 
 
-def get_loss_function(*args, **kwargs):
-    raise NotImplementedError
-
-
 class BCELoss(nn.Module):
     def __init__(self) -> None:
         super().__init__()
@@ -144,7 +140,9 @@ class YOLOLoss:
         # Batch_Size x (Anchor + Class) x H x W
         # TODO: check datatype, why targets has a little bit error with origin version
         predicts, predicts_anc = self.parse_predicts(predicts[0])
-        targets = self.parse_targets(targets, batch_size=predicts.size(0))
+        # TODO: Refactor this operator
+        # targets = self.parse_targets(targets, batch_size=predicts.size(0))
+        targets[:, :, 1:] = targets[:, :, 1:] * self.scale_up
 
         align_targets, valid_masks = self.matcher(targets, predicts)
         # calculate loss between with instance and predict
@@ -162,5 +160,11 @@ class YOLOLoss:
         ## -- DFL -- ##
         loss_dfl = self.dfl(predicts_anc, targets_bbox, valid_masks, box_norm, cls_norm)
 
-        logger.info("Loss IoU: {:.5f}, DFL: {:.5f}, CLS: {:.5f}", loss_iou, loss_dfl, loss_cls)
-        return loss_iou, loss_dfl, loss_cls
+        loss_sum = loss_iou * 0.5 + loss_dfl * 1.5 + loss_cls * 0.5
+        return loss_sum, (loss_iou, loss_dfl, loss_cls)
+
+
+def get_loss_function(cfg: Config) -> YOLOLoss:
+    loss_function = YOLOLoss(cfg)
+    logger.info("✅ Success load loss function")
+    return loss_function