Refractor: Merge Conv Module and Channel Attention Module (#219)

* Fix: Remove some reimported packages

* Refractor: Merge Conv and DepthwiseConv Module into one

* Refractor: Merge Channel Attention

Author: MILK-BIOS

sscma/models/base/general.py

@@ -78,6 +78,7 @@ def __init__(
         conv_layer: Optional[Callable[..., nn.Module]] or Dict or AnyStr = None,
         dilation: int = 1,
         inplace: bool = True,
+        use_depthwise: bool = False,
     ) -> None:
         super().__init__()
         if padding is None:
@@ -86,17 +87,41 @@ def __init__(
             conv_layer = nn.Conv2d
         else:
             conv_layer = get_conv(conv_layer)
-        conv = conv_layer(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride,
-            padding,
-            dilation=dilation,
-            groups=groups,
-            bias=norm_layer is None if bias is None else bias,
-        )
-        self.add_module('conv', conv)
+        if use_depthwise:
+            dw_conv = conv_layer(
+                in_channels,
+                in_channels,
+                kernel_size,
+                stride,
+                padding,
+                dilation=dilation,
+                groups=in_channels,
+                bias=norm_layer is None if bias is None else bias,
+            )
+            pw_conv = conv_layer(
+                in_channels,
+                out_channels,
+                1,
+                stride,
+                padding,
+                dilation=dilation,
+                groups=1,
+                bias=norm_layer is None if bias is None else bias,
+            )
+            self.add_module('dw_conv', dw_conv)
+            self.add_module('pw_conv', pw_conv)
+        else:
+            conv = conv_layer(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride,
+                padding,
+                dilation=dilation,
+                groups=groups,
+                bias=norm_layer is None if bias is None else bias,
+            )
+            self.add_module('conv', conv)
         if norm_layer is not None:
             norm_layer = get_norm(norm_layer)
             self.add_module('norm', norm_layer(out_channels))

sscma/models/detectors/base.py

@@ -1,25 +1,25 @@
 # Copyright (c) Seeed Technology Co.,Ltd. All rights reserved.
-from typing import Dict, Optional, Union, Tuple, List
-from abc import ABCMeta, abstractmethod
 import copy
+from abc import ABCMeta, abstractmethod
+from typing import Dict, List, Optional, Tuple, Union
 
+import torch
 from mmdet.models.detectors import BaseDetector, SemiBaseDetector
-from mmdet.structures import DetDataSample, OptSampleList, SampleList
-from mmdet.utils import OptConfigType, OptMultiConfig, ConfigType, InstanceList
 from mmdet.models.utils import rename_loss_dict, reweight_loss_dict
-from mmdet.structures import SampleList
+from mmdet.structures import DetDataSample, OptSampleList, SampleList
+from mmdet.utils import ConfigType, InstanceList, OptConfigType, OptMultiConfig
 from mmengine.model import BaseModel
-import torch
-from torch import Tensor
 from mmengine.optim import OptimWrapper
-from ..utils import samplelist_boxtype2tensor
+from torch import Tensor
 
-from sscma.registry import MODELS
 from sscma.models.semi import BasePseudoLabelCreator
+from sscma.registry import MODELS
 
+from ..utils import samplelist_boxtype2tensor
 
 ForwardResults = Union[Dict[str, torch.Tensor], List[DetDataSample], Tuple[torch.Tensor], torch.Tensor]
 
+
 @MODELS.register_module()
 class BaseSsod(SemiBaseDetector):
     teacher: BaseDetector

sscma/models/layers/csp_layer.py

@@ -2,81 +2,60 @@
 # Copyright (c) OpenMMLab.
 import torch
 import torch.nn as nn
-from sscma.models.base import ConvModule, DepthwiseSeparableConvModule
+from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig
 from mmengine.model import BaseModule
 from torch import Tensor
 
-from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig
+from sscma.models.base import ConvNormActivation
 
-
-class ChannelAttention(BaseModule):
-    def __init__(self, channels: int, init_cfg: OptMultiConfig = None) -> None:
-        super().__init__(init_cfg=init_cfg)
-        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
-        self.fc = nn.Conv2d(channels, channels, 1, 1, 0, bias=True)
-        self.act = nn.Hardsigmoid(inplace=True)
-
-    def forward(self, x: Tensor) -> Tensor:
-        """Forward function for ChannelAttention."""
-        with torch.cuda.amp.autocast(enabled=False):
-            out = self.global_avgpool(x)
-        out = self.fc(out)
-        out = self.act(out)
-        return x * out
+from .attention import ChannelAttention
 
 
 class CSPLayer(BaseModule):
-    def __init__(self,
-                 in_channels: int,
-                 out_channels: int,
-                 expand_ratio: float = 0.5,
-                 num_blocks: int = 1,
-                 add_identity: bool = True,
-                 use_depthwise: bool = False,
-                 use_cspnext_block: bool = False,
-                 channel_attention: bool = False,
-                 conv_cfg: OptConfigType = None,
-                 norm_cfg: ConfigType = dict(
-                     type='BN', momentum=0.03, eps=0.001),
-                 act_cfg: ConfigType = dict(type='Swish'),
-                 init_cfg: OptMultiConfig = None) -> None:
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        expand_ratio: float = 0.5,
+        num_blocks: int = 1,
+        add_identity: bool = True,
+        use_depthwise: bool = False,
+        use_cspnext_block: bool = False,
+        channel_attention: bool = False,
+        conv_cfg: OptConfigType = None,
+        norm_cfg: ConfigType = dict(type='BN', momentum=0.03, eps=0.001),
+        act_cfg: ConfigType = dict(type='Swish'),
+        init_cfg: OptMultiConfig = None,
+    ) -> None:
         super().__init__(init_cfg=init_cfg)
         block = CSPNeXtBlock if use_cspnext_block else DarknetBottleneck
         mid_channels = int(out_channels * expand_ratio)
         self.channel_attention = channel_attention
-        self.main_conv = ConvModule(
-            in_channels,
-            mid_channels,
-            1,
-            conv_cfg=conv_cfg,
-            norm_cfg=norm_cfg,
-            act_cfg=act_cfg)
-        self.short_conv = ConvModule(
-            in_channels,
-            mid_channels,
-            1,
-            conv_cfg=conv_cfg,
-            norm_cfg=norm_cfg,
-            act_cfg=act_cfg)
-        self.final_conv = ConvModule(
-            2 * mid_channels,
-            out_channels,
-            1,
-            conv_cfg=conv_cfg,
-            norm_cfg=norm_cfg,
-            act_cfg=act_cfg)
-
-        self.blocks = nn.Sequential(*[
-            block(
-                mid_channels,
-                mid_channels,
-                1.0,
-                add_identity,
-                use_depthwise,
-                conv_cfg=conv_cfg,
-                norm_cfg=norm_cfg,
-                act_cfg=act_cfg) for _ in range(num_blocks)
-        ])
+        self.main_conv = ConvNormActivation(
+            in_channels, mid_channels, 1, conv_layer=conv_cfg, norm_layer=norm_cfg, activation_layer=act_cfg
+        )
+        self.short_conv = ConvNormActivation(
+            in_channels, mid_channels, 1, conv_layer=conv_cfg, norm_layer=norm_cfg, activation_layer=act_cfg
+        )
+        self.final_conv = ConvNormActivation(
+            2 * mid_channels, out_channels, 1, conv_layer=conv_cfg, norm_layer=norm_cfg, activation_layer=act_cfg
+        )
+
+        self.blocks = nn.Sequential(
+            *[
+                block(
+                    mid_channels,
+                    mid_channels,
+                    1.0,
+                    add_identity,
+                    use_depthwise,
+                    conv_cfg=conv_cfg,
+                    norm_cfg=norm_cfg,
+                    act_cfg=act_cfg,
+                )
+                for _ in range(num_blocks)
+            ]
+        )
         if channel_attention:
             self.attention = ChannelAttention(2 * mid_channels)
 
@@ -92,40 +71,44 @@ def forward(self, x: Tensor) -> Tensor:
         if self.channel_attention:
             x_final = self.attention(x_final)
         return self.final_conv(x_final)
-    
+
+
 class DarknetBottleneck(BaseModule):
-    def __init__(self,
-                 in_channels: int,
-                 out_channels: int,
-                 expansion: float = 0.5,
-                 add_identity: bool = True,
-                 use_depthwise: bool = False,
-                 conv_cfg: OptConfigType = None,
-                 norm_cfg: ConfigType = dict(
-                     type='BN', momentum=0.03, eps=0.001),
-                 act_cfg: ConfigType = dict(type='Swish'),
-                 init_cfg: OptMultiConfig = None) -> None:
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        expansion: float = 0.5,
+        add_identity: bool = True,
+        use_depthwise: bool = False,
+        conv_cfg: OptConfigType = None,
+        norm_cfg: ConfigType = dict(type='BN', momentum=0.03, eps=0.001),
+        act_cfg: ConfigType = dict(type='Swish'),
+        init_cfg: OptMultiConfig = None,
+    ) -> None:
         super().__init__(init_cfg=init_cfg)
         hidden_channels = int(out_channels * expansion)
-        conv = DepthwiseSeparableConvModule if use_depthwise else ConvModule
-        self.conv1 = ConvModule(
+        self.conv1 = ConvNormActivation(
             in_channels,
             hidden_channels,
             1,
-            conv_cfg=conv_cfg,
-            norm_cfg=norm_cfg,
-            act_cfg=act_cfg)
-        self.conv2 = conv(
+            conv_layer=conv_cfg,
+            norm_layer=norm_cfg,
+            activation_layer=act_cfg,
+            use_depthwise=False,
+        )
+        self.conv2 = ConvNormActivation(
             hidden_channels,
             out_channels,
             3,
             stride=1,
             padding=1,
-            conv_cfg=conv_cfg,
-            norm_cfg=norm_cfg,
-            act_cfg=act_cfg)
-        self.add_identity = \
-            add_identity and in_channels == out_channels
+            conv_layer=conv_cfg,
+            norm_layer=norm_cfg,
+            activation_layer=act_cfg,
+            use_depthwise=use_depthwise,
+        )
+        self.add_identity = add_identity and in_channels == out_channels
 
     def forward(self, x: Tensor) -> Tensor:
         """Forward function."""
@@ -140,40 +123,43 @@ def forward(self, x: Tensor) -> Tensor:
 
 
 class CSPNeXtBlock(BaseModule):
-    def __init__(self,
-                 in_channels: int,
-                 out_channels: int,
-                 expansion: float = 0.5,
-                 add_identity: bool = True,
-                 use_depthwise: bool = False,
-                 kernel_size: int = 5,
-                 conv_cfg: OptConfigType = None,
-                 norm_cfg: ConfigType = dict(
-                     type='BN', momentum=0.03, eps=0.001),
-                 act_cfg: ConfigType = dict(type='SiLU'),
-                 init_cfg: OptMultiConfig = None) -> None:
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        expansion: float = 0.5,
+        add_identity: bool = True,
+        use_depthwise: bool = False,
+        kernel_size: int = 5,
+        conv_cfg: OptConfigType = None,
+        norm_cfg: ConfigType = dict(type='BN', momentum=0.03, eps=0.001),
+        act_cfg: ConfigType = dict(type='SiLU'),
+        init_cfg: OptMultiConfig = None,
+    ) -> None:
         super().__init__(init_cfg=init_cfg)
         hidden_channels = int(out_channels * expansion)
-        conv = DepthwiseSeparableConvModule if use_depthwise else ConvModule
-        self.conv1 = conv(
+        self.conv1 = ConvNormActivation(
             in_channels,
             hidden_channels,
             3,
             stride=1,
             padding=1,
-            norm_cfg=norm_cfg,
-            act_cfg=act_cfg)
-        self.conv2 = DepthwiseSeparableConvModule(
+            norm_layer=norm_cfg,
+            activation_layer=act_cfg,
+            use_depthwise=use_depthwise,
+        )
+        self.conv2 = ConvNormActivation(
             hidden_channels,
             out_channels,
             kernel_size,
             stride=1,
             padding=kernel_size // 2,
-            conv_cfg=conv_cfg,
-            norm_cfg=norm_cfg,
-            act_cfg=act_cfg)
-        self.add_identity = \
-            add_identity and in_channels == out_channels
+            conv_layer=conv_cfg,
+            norm_layer=norm_cfg,
+            activation_layer=act_cfg,
+            use_depthwise=True,
+        )
+        self.add_identity = add_identity and in_channels == out_channels
 
     def forward(self, x: Tensor) -> Tensor:
         """Forward function."""