Merge branch 'MODEL' of github.com:LucyTuan/yolov9mit into MODEL

yolo/model/module.py

@@ -1,10 +1,10 @@
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import torch
 from torch import Tensor, nn
 from torch.nn.common_types import _size_2_t
 
-from yolo.tools.module_helper import auto_pad, get_activation
+from yolo.tools.module_helper import auto_pad, get_activation, round_up
 
 
 class Conv(nn.Module):
@@ -99,6 +99,47 @@ class SPPELAN(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):
+        super().__init__()
+
+        groups = 4 if use_group else 1
+        anchor_channels = 4 * reg_max
+        # 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))
+
+        self.anchor_conv = nn.Sequential(
+            Conv(in_channels, anchor_neck, 3),
+            Conv(anchor_neck, anchor_neck, 3, groups=groups),
+            nn.Conv2d(anchor_neck, anchor_channels, 1, groups=groups),
+        )
+        self.class_conv = nn.Sequential(
+            Conv(in_channels, class_neck, 3), Conv(class_neck, class_neck, 3), nn.Conv2d(class_neck, num_classes, 1)
+        )
+
+    def forward(self, x: List[Tensor]) -> List[Tensor]:
+        anchor_x = self.anchor_conv(x)
+        class_x = self.class_conv(x)
+        return torch.cat([anchor_x, class_x], dim=1)
+
+
+class MultiheadDetection(nn.Module):
+    """Mutlihead Detection module for Dual detect or Triple detect"""
+
+    def __init__(self, in_channels: List[int], num_classes: int, **head_kwargs):
+        super().__init__()
+        self.heads = nn.ModuleList(
+            [Detection(head_in_channels, num_classes, **head_kwargs) for head_in_channels in in_channels]
+        )
+
+    def forward(self, x_list: List[torch.Tensor]) -> List[torch.Tensor]:
+        return [head(x) for x, head in zip(x_list, self.heads)]
+
+
+#### -- ####
 # RepVGG
 class RepConv(nn.Module):
     """A convolutional block that combines two convolution layers (kernel and point-wise)."""

yolo/model/yolo.py

@@ -27,8 +27,9 @@ class YOLO(nn.Module):
         model_list = nn.ModuleList()
         output_dim = [3]
         layer_indices_by_tag = {}
+        logger.info(f"🚜 Building YOLO")
         for arch_name in model_arch:
-            logger.info(f"🏗️  Building model-{arch_name}")
+            logger.info(f"  🏗️  Building {arch_name}")
             for layer_idx, layer_spec in enumerate(model_arch[arch_name], start=1):
                 layer_type, layer_info = next(iter(layer_spec.items()))
                 layer_args = layer_info.get("args", {})
@@ -102,9 +103,3 @@ def get_model(model_cfg: dict) -> YOLO:
     model = YOLO(model_cfg)
     logger.info("✅ Success load model")
     return model
-
-
-if __name__ == "__main__":
-    model_cfg = load_model_cfg("v7-base")
-
-    YOLO(model_cfg)

yolo/tools/module_helper.py

@@ -1,6 +1,6 @@
-from typing import Tuple
+from typing import Tuple, Union
 
-from torch import nn
+from torch import Tensor, nn
 from torch.nn.common_types import _size_2_t
 
 
@@ -34,3 +34,27 @@ def get_activation(activation: str) -> nn.Module:
         return activation_map[activation.lower()]()
     else:
         raise ValueError(f"Activation function '{activation}' is not found in torch.nn")
+
+
+def round_up(x: Union[int, Tensor], div: int = 1) -> Union[int, Tensor]:
+    """
+    Rounds up `x` to the bigger-nearest multiple of `div`.
+    """
+    return x + (-x % div)
+
+
+def make_chunk(input_list, chunk_num):
+    """
+    Args: input_list: [0, 1, 2, 3, 4, 5], chunk: 2
+    Return: [[0, 1, 2], [3, 4, 5]]
+    """
+    list_size = len(input_list)
+
+    if list_size % chunk_num != 0:
+        raise ValueError(
+            f"The length of the input list ({list_size}) must be exactly\
+                            divisible by the number of chunks ({chunk_num})."
+        )
+
+    chunk_size = list_size // chunk_num
+    return [input_list[i : i + chunk_size] for i in range(0, list_size, chunk_size)]