# Copyright (c) OpenMMLab. All rights reserved. import torch import torch.nn as nn from mmcv.cnn import ConvModule, DepthwiseSeparableConvModule from mmengine.model import BaseModule from torch import Tensor from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig from .se_layer import ChannelAttention class DarknetBottleneck(BaseModule): """The basic bottleneck block used in Darknet. Each ResBlock consists of two ConvModules and the input is added to the final output. Each ConvModule is composed of Conv, BN, and LeakyReLU. The first convLayer has filter size of 1x1 and the second one has the filter size of 3x3. Args: in_channels (int): The input channels of this Module. out_channels (int): The output channels of this Module. expansion (float): The kernel size of the convolution. Defaults to 0.5. add_identity (bool): Whether to add identity to the out. Defaults to True. use_depthwise (bool): Whether to use depthwise separable convolution. Defaults to False. conv_cfg (dict): Config dict for convolution layer. Defaults to None, which means using conv2d. norm_cfg (dict): Config dict for normalization layer. Defaults to dict(type='BN'). act_cfg (dict): Config dict for activation layer. Defaults to dict(type='Swish'). """ 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( in_channels, hidden_channels, 1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg) self.conv2 = conv( 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 def forward(self, x: Tensor) -> Tensor: """Forward function.""" identity = x out = self.conv1(x) out = self.conv2(out) if self.add_identity: return out + identity else: return out class CSPNeXtBlock(BaseModule): """The basic bottleneck block used in CSPNeXt. Args: in_channels (int): The input channels of this Module. out_channels (int): The output channels of this Module. expansion (float): Expand ratio of the hidden channel. Defaults to 0.5. add_identity (bool): Whether to add identity to the out. Only works when in_channels == out_channels. Defaults to True. use_depthwise (bool): Whether to use depthwise separable convolution. Defaults to False. kernel_size (int): The kernel size of the second convolution layer. Defaults to 5. conv_cfg (dict): Config dict for convolution layer. Defaults to None, which means using conv2d. norm_cfg (dict): Config dict for normalization layer. Defaults to dict(type='BN', momentum=0.03, eps=0.001). act_cfg (dict): Config dict for activation layer. Defaults to dict(type='SiLU'). init_cfg (:obj:`ConfigDict` or dict or list[dict] or list[:obj:`ConfigDict`], optional): Initialization config dict. Defaults to 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( in_channels, hidden_channels, 3, stride=1, padding=1, norm_cfg=norm_cfg, act_cfg=act_cfg) self.conv2 = DepthwiseSeparableConvModule( 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 def forward(self, x: Tensor) -> Tensor: """Forward function.""" identity = x out = self.conv1(x) out = self.conv2(out) if self.add_identity: return out + identity else: return out class CSPLayer(BaseModule): """Cross Stage Partial Layer. Args: in_channels (int): The input channels of the CSP layer. out_channels (int): The output channels of the CSP layer. expand_ratio (float): Ratio to adjust the number of channels of the hidden layer. Defaults to 0.5. num_blocks (int): Number of blocks. Defaults to 1. add_identity (bool): Whether to add identity in blocks. Defaults to True. use_cspnext_block (bool): Whether to use CSPNeXt block. Defaults to False. use_depthwise (bool): Whether to use depthwise separable convolution in blocks. Defaults to False. channel_attention (bool): Whether to add channel attention in each stage. Defaults to True. conv_cfg (dict, optional): Config dict for convolution layer. Defaults to None, which means using conv2d. norm_cfg (dict): Config dict for normalization layer. Defaults to dict(type='BN') act_cfg (dict): Config dict for activation layer. Defaults to dict(type='Swish') init_cfg (:obj:`ConfigDict` or dict or list[dict] or list[:obj:`ConfigDict`], optional): Initialization config dict. Defaults to 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) ]) if channel_attention: self.attention = ChannelAttention(2 * mid_channels) def forward(self, x: Tensor) -> Tensor: """Forward function.""" x_short = self.short_conv(x) x_main = self.main_conv(x) x_main = self.blocks(x_main) x_final = torch.cat((x_main, x_short), dim=1) if self.channel_attention: x_final = self.attention(x_final) return self.final_conv(x_final)