[!HOTFIX] Add use_identity block to load rt-detr pre-trained weight…

… correctly
Nota-NetsPresso · Nov 28, 2024 · 3313098 · 3313098
1 parent 45ce412
commit 3313098
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/src/netspresso_trainer/models/necks/experimental/rtdetr_hybrid_encoder.py b/src/netspresso_trainer/models/necks/experimental/rtdetr_hybrid_encoder.py
@@ -154,7 +154,7 @@ def __init__(
         for _ in range(len(self.in_channels) - 1, 0, -1):
             self.lateral_convs.append(ConvLayer(self.hidden_dim, self.hidden_dim, kernel_size=1, stride=1, act_type=act))
             self.fpn_blocks.append(
-                CSPRepLayer(self.hidden_dim * 2, self.hidden_dim, round(3 * depth_mult), act=act, expansion=expansion)
+                CSPRepLayer(self.hidden_dim * 2, self.hidden_dim, round(3 * depth_mult), act=act, expansion=expansion, use_identity=False)
             )
 
         # bottom-up pan
@@ -165,7 +165,7 @@ def __init__(
                 ConvLayer(self.hidden_dim, self.hidden_dim, kernel_size=3, stride=2, act_type=act)
             )
             self.pan_blocks.append(
-                CSPRepLayer(self.hidden_dim * 2, self.hidden_dim, round(3 * depth_mult), act=act, expansion=expansion)
+                CSPRepLayer(self.hidden_dim * 2, self.hidden_dim, round(3 * depth_mult), act=act, expansion=expansion, use_identity=False)
             )
         self._reset_parameters()
 

diff --git a/src/netspresso_trainer/models/op/custom.py b/src/netspresso_trainer/models/op/custom.py
@@ -195,7 +195,9 @@ def __init__(self,
                  out_channels: int,
                  kernel_size: Union[int, Tuple[int, int]] = 3,
                  groups: int = 1,
-                 act_type: Optional[str] = None,):
+                 act_type: Optional[str] = None,
+                 use_identity: Optional[bool]=True,
+                 ):
         if act_type is None:
             act_type = 'silu'
         super().__init__()
@@ -209,7 +211,7 @@ def __init__(self,
         self.groups = groups
         self.conv1 = ConvLayer(in_channels, out_channels, kernel_size, groups=groups, use_act=False)
         self.conv2 = ConvLayer(in_channels, out_channels, 1, groups=groups, use_act=False)
-        self.rbr_identity = nn.BatchNorm2d(num_features=in_channels) if out_channels == in_channels else nn.Identity()
+        self.rbr_identity = nn.BatchNorm2d(num_features=in_channels) if use_identity and out_channels == in_channels else nn.Identity()
 
         assert act_type in ACTIVATION_REGISTRY
         self.act = ACTIVATION_REGISTRY[act_type]()
@@ -851,13 +853,14 @@ def __init__(self,
                  num_blocks: int=3,
                  expansion: float=1.0,
                  bias: bool= False,
+                 use_identity: Optional[bool]=True,
                  act: str="silu"):
         super(CSPRepLayer, self).__init__()
         hidden_channels = int(out_channels * expansion)
         self.conv1 = ConvLayer(in_channels, hidden_channels, kernel_size=1, stride=1, bias=bias, act_type=act)
         self.conv2 = ConvLayer(in_channels, hidden_channels, kernel_size=1, stride=1, bias=bias, act_type=act)
         self.bottlenecks = nn.Sequential(*[
-            RepVGGBlock(hidden_channels, hidden_channels, act_type=act) for _ in range(num_blocks)
+            RepVGGBlock(hidden_channels, hidden_channels, act_type=act, use_identity=use_identity) for _ in range(num_blocks)
         ])
         if hidden_channels != out_channels:
             self.conv3 = ConvLayer(hidden_channels, out_channels, kernel_size=1, stride=1, bias=bias, act_type=act)