diff --git a/mmdet/models/dense_heads/rtmdet_head.py b/mmdet/models/dense_heads/rtmdet_head.py index 3c53b68669d..ae0ee6d2f35 100644 --- a/mmdet/models/dense_heads/rtmdet_head.py +++ b/mmdet/models/dense_heads/rtmdet_head.py @@ -3,7 +3,7 @@ import torch import torch.nn as nn -from mmcv.cnn import ConvModule, Scale, is_norm +from mmcv.cnn import ConvModule, DepthwiseSeparableConvModule, Scale, is_norm from mmengine.model import bias_init_with_prob, constant_init, normal_init from mmengine.structures import InstanceData from torch import Tensor @@ -536,6 +536,8 @@ class RTMDetSepBNHead(RTMDetHead): in_channels (int): Number of channels in the input feature map. share_conv (bool): Whether to share conv layers between stages. Defaults to True. + use_depthwise (bool): Whether to use depthwise separable convolution in + head. Defaults to False. norm_cfg (:obj:`ConfigDict` or dict)): Config dict for normalization layer. Defaults to dict(type='BN', momentum=0.03, eps=0.001). act_cfg (:obj:`ConfigDict` or dict)): Config dict for activation layer. @@ -547,6 +549,7 @@ def __init__(self, num_classes: int, in_channels: int, share_conv: bool = True, + use_depthwise: bool = False, norm_cfg: ConfigType = dict( type='BN', momentum=0.03, eps=0.001), act_cfg: ConfigType = dict(type='SiLU'), @@ -555,6 +558,7 @@ def __init__(self, **kwargs) -> None: self.share_conv = share_conv self.exp_on_reg = exp_on_reg + self.use_depthwise = use_depthwise super().__init__( num_classes, in_channels, @@ -565,6 +569,8 @@ def __init__(self, def _init_layers(self) -> None: """Initialize layers of the head.""" + conv = DepthwiseSeparableConvModule \ + if self.use_depthwise else ConvModule self.cls_convs = nn.ModuleList() self.reg_convs = nn.ModuleList() @@ -578,7 +584,7 @@ def _init_layers(self) -> None: for i in range(self.stacked_convs): chn = self.in_channels if i == 0 else self.feat_channels cls_convs.append( - ConvModule( + conv( chn, self.feat_channels, 3, @@ -588,7 +594,7 @@ def _init_layers(self) -> None: norm_cfg=self.norm_cfg, act_cfg=self.act_cfg)) reg_convs.append( - ConvModule( + conv( chn, self.feat_channels, 3, diff --git a/mmdet/models/necks/cspnext_pafpn.py b/mmdet/models/necks/cspnext_pafpn.py index 0001514db64..a52ba72d9b3 100644 --- a/mmdet/models/necks/cspnext_pafpn.py +++ b/mmdet/models/necks/cspnext_pafpn.py @@ -82,6 +82,7 @@ def __init__( in_channels[idx - 1], num_blocks=num_csp_blocks, add_identity=False, + use_depthwise=use_depthwise, use_cspnext_block=True, expand_ratio=expand_ratio, conv_cfg=conv_cfg, @@ -108,6 +109,7 @@ def __init__( in_channels[idx + 1], num_blocks=num_csp_blocks, add_identity=False, + use_depthwise=use_depthwise, use_cspnext_block=True, expand_ratio=expand_ratio, conv_cfg=conv_cfg, @@ -117,7 +119,7 @@ def __init__( self.out_convs = nn.ModuleList() for i in range(len(in_channels)): self.out_convs.append( - ConvModule( + conv( in_channels[i], out_channels, 3,