🔀 [Merge] branch 'MODEL' into TEST

yolo/config/config.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
 from typing import Dict, List, Union
 
+from torch import nn
+
 
 @dataclass
 class AnchorConfig:
@@ -100,6 +102,17 @@ class Download:
     datasets: Datasets
 
 
+@dataclass
+class YOLOLayer(nn.Module):
+    source: Union[int, str, List[int]]
+    output: bool
+    tags: str
+    layer_type: str
+
+    def __post_init__(self):
+        super().__init__()
+
+
 @dataclass
 class Config:
     model: Model

yolo/model/module.py

@@ -24,7 +24,7 @@ class Conv(nn.Module):
     ):
         super().__init__()
         kwargs.setdefault("padding", auto_pad(kernel_size, **kwargs))
-        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, **kwargs)
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias=False, **kwargs)
         self.bn = nn.BatchNorm2d(out_channels)
         self.act = get_activation(activation)
 
@@ -49,14 +49,16 @@ class Pool(nn.Module):
 class Detection(nn.Module):
     """A single YOLO Detection head for detection models"""
 
-    def __init__(self, in_channels: int, num_classes: int, *, reg_max: int = 16, use_group: bool = True):
+    def __init__(self, in_channels: Tuple[int], num_classes: int, *, reg_max: int = 16, use_group: bool = True):
         super().__init__()
 
         groups = 4 if use_group else 1
         anchor_channels = 4 * reg_max
+
+        first_neck, in_channels = in_channels
         # TODO: round up head[0] channels or each head?
-        anchor_neck = max(round_up(in_channels // 4, groups), anchor_channels, 16)
-        class_neck = max(in_channels, min(num_classes * 2, 128))
+        anchor_neck = max(round_up(first_neck // 4, groups), anchor_channels, 16)
+        class_neck = max(first_neck, min(num_classes * 2, 128))
 
         self.anchor_conv = nn.Sequential(
             Conv(in_channels, anchor_neck, 3),
@@ -78,8 +80,12 @@ class MultiheadDetection(nn.Module):
 
     def __init__(self, in_channels: List[int], num_classes: int, **head_kwargs):
         super().__init__()
+        # TODO: Refactor these parts
         self.heads = nn.ModuleList(
-            [Detection(head_in_channels, num_classes, **head_kwargs) for head_in_channels in in_channels]
+            [
+                Detection((in_channels[3 * (idx // 3)], in_channel), num_classes, **head_kwargs)
+                for idx, in_channel in enumerate(in_channels)
+            ]
         )
 
     def forward(self, x_list: List[torch.Tensor]) -> List[torch.Tensor]:
@@ -118,7 +124,7 @@ class RepNBottleneck(nn.Module):
         *,
         kernel_size: Tuple[int, int] = (3, 3),
         residual: bool = True,
-        expand: float = 0.5,
+        expand: float = 1.0,
         **kwargs
     ):
         super().__init__()

yolo/model/yolo.py

@@ -4,8 +4,9 @@ import torch.nn as nn
 from loguru import logger
 from omegaconf import ListConfig, OmegaConf
 
-from yolo.config.config import Config, Model
+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):
@@ -21,13 +22,13 @@ class YOLO(nn.Module):
         super(YOLO, self).__init__()
         self.num_classes = num_classes
         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]]]]):
-        model_list = nn.ModuleList()
-        output_dim = [3]
-        layer_indices_by_tag = {}
-        layer_idx = 1
+        self.layer_index = {}
+        output_dim, layer_idx = [3], 1
         logger.info(f"🚜 Building YOLO")
         for arch_name in model_arch:
             logger.info(f"  🏗️  Building {arch_name}")
@@ -36,11 +37,7 @@ class YOLO(nn.Module):
                 layer_args = layer_info.get("args", {})
 
                 # Get input source
-                source = layer_info.get("source", -1)
-                if isinstance(source, str):
-                    source = layer_indices_by_tag[source]
-                elif isinstance(source, ListConfig):
-                    source = [layer_indices_by_tag[idx] if isinstance(idx, str) else idx for idx in source]
+                source = self.get_source_idx(layer_info.get("source", -1), layer_idx)
 
                 # Find in channels
                 if any(module in layer_type for module in ["Conv", "ELAN", "ADown", "CBLinear"]):
@@ -51,29 +48,29 @@ class YOLO(nn.Module):
 
                 # create layers
                 layer = self.create_layer(layer_type, source, layer_info, **layer_args)
-                model_list.append(layer)
+                self.model.append(layer)
 
-                if "tags" in layer_info:
-                    if layer_info["tags"] in layer_indices_by_tag:
+                if layer.tags:
+                    if layer.tags in self.layer_index:
                         raise ValueError(f"Duplicate tag '{layer_info['tags']}' found.")
-                    layer_indices_by_tag[layer_info["tags"]] = layer_idx
+                    self.layer_index[layer.tags] = layer_idx
 
                 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
 
-        self.model = model_list
-
     def forward(self, x):
-        y = [x]
+        y = {0: x}
         output = []
-        for layer in self.model:
+        for index, layer in enumerate(self.model, start=1):
             if isinstance(layer.source, list):
                 model_input = [y[idx] for idx in layer.source]
             else:
                 model_input = y[layer.source]
             x = layer(model_input)
-            y.append(x)
+            if hasattr(layer, "save"):
+                y[index] = x
             if layer.output:
                 output.append(x)
         return output
@@ -90,10 +87,23 @@ class YOLO(nn.Module):
         if layer_type == "IDetect":
             return None
 
-    def create_layer(self, layer_type: str, source: Union[int, list], layer_info, **kwargs):
+    def get_source_idx(self, source: Union[ListConfig, str, int], layer_idx: int):
+        if isinstance(source, ListConfig):
+            return [self.get_source_idx(index, layer_idx) for index in source]
+        if isinstance(source, str):
+            source = self.layer_index[source]
+        if source < 0:
+            source += layer_idx
+        if source > 0:
+            setattr(self.model[source - 1], "save", True)
+        return source
+
+    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)

yolo/tools/module_helper.py

@@ -31,7 +31,7 @@ def get_activation(activation: str) -> nn.Module:
         if isinstance(obj, type) and issubclass(obj, nn.Module)
     }
     if activation.lower() in activation_map:
-        return activation_map[activation.lower()]()
+        return activation_map[activation.lower()](inplace=True)
     else:
         raise ValueError(f"Activation function '{activation}' is not found in torch.nn")