✨ [Add] model logger to print model as table

yolo/model/yolo.py

@@ -6,6 +6,7 @@ from omegaconf import ListConfig, OmegaConf
 
 from yolo.config.config import Config, Model, YOLOLayer
 from yolo.tools.layer_helper import get_layer_map
+from yolo.tools.log_helper import log_model
 
 
 class YOLO(nn.Module):
@@ -23,6 +24,7 @@ class YOLO(nn.Module):
         self.layer_map = get_layer_map()  # Get the map Dict[str: Module]
         self.model: List[YOLOLayer] = nn.ModuleList()
         self.build_model(model_cfg.model)
+        log_model(self.model)
 
     def build_model(self, model_arch: Dict[str, List[Dict[str, Dict[str, Dict]]]]):
         self.layer_index = {}
@@ -55,6 +57,7 @@ class YOLO(nn.Module):
 
                 out_channels = self.get_out_channels(layer_type, layer_args, output_dim, source)
                 output_dim.append(out_channels)
+                setattr(layer, "out_c", out_channels)
             layer_idx += 1
 
     def forward(self, x):
@@ -98,7 +101,9 @@ class YOLO(nn.Module):
     def create_layer(self, layer_type: str, source: Union[int, list], layer_info: Dict, **kwargs) -> YOLOLayer:
         if layer_type in self.layer_map:
             layer = self.layer_map[layer_type](**kwargs)
+            setattr(layer, "layer_type", layer_type)
             setattr(layer, "source", source)
+            setattr(layer, "in_c", kwargs.get("in_channels", None))
             setattr(layer, "output", layer_info.get("output", False))
             setattr(layer, "tags", layer_info.get("tags", None))
             return layer

yolo/tools/log_helper.py

@@ -12,8 +12,13 @@ Example:
 """
 
 import sys
+from typing import List
 
 from loguru import logger
+from rich.console import Console
+from rich.table import Table
+
+from yolo.config.config import YOLOLayer
 
 
 def custom_logger():
@@ -22,3 +27,24 @@ def custom_logger():
         sys.stderr,
         format="<green>{time:MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>",
     )
+
+
+def log_model(model: List[YOLOLayer]):
+    console = Console()
+    table = Table(title="Model Layers")
+
+    table.add_column("Index", justify="center")
+    table.add_column("Layer Type", justify="center")
+    table.add_column("Tags", justify="center")
+    table.add_column("Params", justify="right")
+    table.add_column("Channels (IN->OUT)", justify="center")
+
+    for idx, layer in enumerate(model, start=1):
+        layer_param = sum(x.numel() for x in layer.parameters())  # number parameters
+        in_channels, out_channels = getattr(layer, "in_c", None), getattr(layer, "out_c", None)
+        if in_channels and out_channels:
+            channels = f"{in_channels:4} -> {out_channels:4}"
+        else:
+            channels = "-"
+        table.add_row(str(idx), layer.layer_type, layer.tags, f"{layer_param:,}", channels)
+    console.print(table)