🔀 [Merge] branch 'MODEL' into TEST

config/model/v7-base.yaml

@@ -241,3 +241,4 @@ model:
             - [36,75, 76,55, 72,146]  # P4/16
             - [142,110, 192,243, 459,401]  # P5/32
       source: [102, 103, 104]
+      output: True

model/module.py

@@ -11,10 +11,10 @@ class Conv(nn.Module):
         out_channels,
         kernel_size,
         stride=1,
-        padding=0,
+        padding=None,
         dilation=1,
         groups=1,
-        act=nn.ReLU(),
+        act=nn.SiLU(),
         bias=False,
         auto_padding=True,
         padding_mode="zeros",
@@ -48,10 +48,10 @@ class Conv(nn.Module):
 # RepVGG
 class RepConv(nn.Module):
     # https://github.com/DingXiaoH/RepVGG
-    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, groups=1, act=nn.ReLU()):
+    def __init__(self, in_channels, out_channels, kernel_size=3, padding=None, stride=1, groups=1, act=nn.SiLU(), deploy=False):
 
         super().__init__()
-
+        self.deploy = deploy 
         self.conv1 = Conv(in_channels, out_channels, kernel_size, stride, groups=groups, act=False)
         self.conv2 = Conv(in_channels, out_channels, 1, stride, groups=groups, act=False)
         self.act = act if isinstance(act, nn.Module) else nn.Identity()
@@ -64,6 +64,28 @@ class RepConv(nn.Module):
 
     # to be implement
     # def fuse_convs(self):
+    def fuse_conv_bn(self, conv, bn):
+
+        std = (bn.running_var + bn.eps).sqrt()
+        bias = bn.bias - bn.running_mean * bn.weight / std
+
+        t = (bn.weight / std).reshape(-1, 1, 1, 1)
+        weights = conv.weight * t
+
+        bn = nn.Identity()
+        conv = nn.Conv2d(in_channels = conv.in_channels,
+                              out_channels = conv.out_channels,
+                              kernel_size = conv.kernel_size,
+                              stride=conv.stride,
+                              padding = conv.padding,
+                              dilation = conv.dilation,
+                              groups = conv.groups,
+                              bias = True,
+                              padding_mode = conv.padding_mode)
+
+        conv.weight = torch.nn.Parameter(weights)
+        conv.bias = torch.nn.Parameter(bias)
+        return conv
 
 
 # ResNet

model/yolo.py

@@ -1,10 +1,12 @@
-import torch.nn as nn
+import inspect
+from typing import Any, Dict, List, Union
+
 import torch
+import torch.nn as nn
 from loguru import logger
-from typing import Dict, Any, List, Union
-import inspect
-from utils.tools import load_model_cfg
+
 from model import module
+from utils.tools import load_model_cfg
 
 
 def get_layer_map():
@@ -46,6 +48,7 @@ class YOLO(nn.Module):
                 layer_type, layer_info = next(iter(layer_spec.items()))
                 layer_args = layer_info.get("args", {})
                 source = layer_info.get("source", -1)
+                output = layer_info.get("output", False)
 
                 if isinstance(source, str):
                     source = layer_indices_by_tag[source]
@@ -55,7 +58,7 @@ class YOLO(nn.Module):
                     layer_args["nc"] = self.nc
                     layer_args["ch"] = [output_dim[idx] for idx in source]
 
-                layer = self.create_layer(layer_type, source, **layer_args)
+                layer = self.create_layer(layer_type, source, output, **layer_args)
                 model_list.append(layer)
 
                 if "tags" in layer_info:
@@ -69,6 +72,7 @@ class YOLO(nn.Module):
 
     def forward(self, x):
         y = [x]
+        output = []
         for layer in self.model:
             if isinstance(layer.source, list):
                 model_input = [y[idx] for idx in layer.source]
@@ -76,7 +80,9 @@ class YOLO(nn.Module):
                 model_input = y[layer.source]
             x = layer(model_input)
             y.append(x)
-        return x
+            if layer.output:
+                output.append(x)
+        return output
 
     def get_out_channels(self, layer_type: str, layer_args: dict, output_dim: list, source: Union[int, list]):
         if "Conv" in layer_type:
@@ -88,10 +94,11 @@ class YOLO(nn.Module):
         if layer_type == "IDetect":
             return None
 
-    def create_layer(self, layer_type: str, source: Union[int, list], **kwargs):
+    def create_layer(self, layer_type: str, source: Union[int, list], output=False, **kwargs):
         if layer_type in self.layer_map:
             layer = self.layer_map[layer_type](**kwargs)
             layer.source = source
+            layer.output = output
             return layer
         else:
             raise ValueError(f"Unsupported layer type: {layer_type}")