💬 [Add] Progress class, handle progress bar

yolo/tools/log_helper.py

@@ -16,6 +16,7 @@ from typing import List
 
 from loguru import logger
 from rich.console import Console
+from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
 from rich.table import Table
 
 from yolo.config.config import YOLOLayer
@@ -29,6 +30,34 @@ def custom_logger():
     )
 
 
+class CustomProgress:
+    def __init__(self):
+        self.progress = Progress(
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(bar_width=None),
+            TextColumn("{task.completed}/{task.total}"),
+            TimeRemainingColumn(),
+        )
+
+    def start_train(self, num_epochs: int):
+        self.task_epoch = self.progress.add_task("[cyan]Epochs", total=num_epochs)
+
+    def one_epoch(self):
+        self.progress.update(self.task_epoch, advance=1)
+
+    def start_batch(self, num_batches):
+        self.batch_task = self.progress.add_task("[green]Batches", total=num_batches)
+
+    def one_batch(self, loss_each):
+        loss_iou, loss_dfl, loss_cls = loss_each
+        # TODO: make it flexible? if need add more loss
+        loss_str = f"Loss IoU: {loss_iou:.3f}, DFL: {loss_dfl:.3f}, CLS: {loss_cls:.3f}"
+        self.progress.update(self.batch_task, advance=1, description=f"[green]Batches {loss_str}")
+
+    def finish_batch(self):
+        self.progress.remove_task(self.batch_task)
+
+
 def log_model(model: List[YOLOLayer]):
     console = Console()
     table = Table(title="Model Layers")

yolo/tools/trainer.py

@@ -1,11 +1,13 @@
 import torch
 from loguru import logger
 from torch import Tensor
+
+# TODO: We may can't use CUDA?
 from torch.cuda.amp import GradScaler, autocast
-from tqdm import tqdm
 
 from yolo.config.config import Config, TrainConfig
 from yolo.model.yolo import YOLO
+from yolo.tools.log_helper import CustomProgress
 from yolo.tools.model_helper import EMA, get_optimizer, get_scheduler
 from yolo.utils.loss import get_loss_function
 
@@ -26,16 +28,13 @@ class Trainer:
             self.ema = None
         self.scaler = GradScaler()
 
-    def train_one_batch(self, data: Tensor, targets: Tensor, progress: tqdm):
+    def train_one_batch(self, data: Tensor, targets: Tensor):
         data, targets = data.to(self.device), targets.to(self.device)
         self.optimizer.zero_grad()
 
         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)
@@ -43,17 +42,21 @@ class Trainer:
 
         return loss.item(), loss_item
 
-        return loss.item()
-
-    def train_one_epoch(self, dataloader):
+    def train_one_epoch(self, dataloader, progress: CustomProgress):
         self.model.train()
         total_loss = 0
-        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()
+        progress.start_batch(len(dataloader))
+
+        for data, targets in dataloader:
+            loss, loss_each = self.train_one_batch(data, targets)
+
+            total_loss += loss
+            progress.one_batch(loss_each)
+
+        if self.scheduler:
+            self.scheduler.step()
+
+        progress.finish_batch()
         return total_loss / len(dataloader)
 
     def save_checkpoint(self, epoch: int, filename="checkpoint.pt"):
@@ -69,9 +72,16 @@ class Trainer:
         torch.save(checkpoint, filename)
 
     def train(self, dataloader, num_epochs):
-        logger.info("start train")
-        for epoch in range(num_epochs):
-            epoch_loss = self.train_one_epoch(dataloader)
-            logger.info(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")
-            if (epoch + 1) % 5 == 0:
-                self.save_checkpoint(epoch, f"checkpoint_epoch_{epoch+1}.pth")
+        logger.info("🚄 Start Training!")
+        progress = CustomProgress()
+
+        with progress.progress:
+            progress.start_train(num_epochs)
+            for epoch in range(num_epochs):
+
+                epoch_loss = self.train_one_epoch(dataloader, progress)
+                progress.one_epoch()
+
+                logger.info(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")
+                if (epoch + 1) % 5 == 0:
+                    self.save_checkpoint(epoch, f"checkpoint_epoch_{epoch+1}.pth")