| |
| import math |
| from typing import Sequence, Tuple |
|
|
| import torch.nn as nn |
| from mmcv.cnn import ConvModule, DepthwiseSeparableConvModule |
| from mmengine.model import BaseModule |
| from torch import Tensor |
| from torch.nn.modules.batchnorm import _BatchNorm |
|
|
| from mmdet.registry import MODELS |
| from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig |
| from ..layers import CSPLayer |
| from .csp_darknet import SPPBottleneck |
|
|
|
|
| @MODELS.register_module() |
| class CSPNeXt(BaseModule): |
| """CSPNeXt backbone used in RTMDet. |
| |
| Args: |
| arch (str): Architecture of CSPNeXt, from {P5, P6}. |
| Defaults to P5. |
| expand_ratio (float): Ratio to adjust the number of channels of the |
| hidden layer. Defaults to 0.5. |
| deepen_factor (float): Depth multiplier, multiply number of |
| blocks in CSP layer by this amount. Defaults to 1.0. |
| widen_factor (float): Width multiplier, multiply number of |
| channels in each layer by this amount. Defaults to 1.0. |
| out_indices (Sequence[int]): Output from which stages. |
| Defaults to (2, 3, 4). |
| frozen_stages (int): Stages to be frozen (stop grad and set eval |
| mode). -1 means not freezing any parameters. Defaults to -1. |
| use_depthwise (bool): Whether to use depthwise separable convolution. |
| Defaults to False. |
| arch_ovewrite (list): Overwrite default arch settings. |
| Defaults to None. |
| spp_kernel_sizes: (tuple[int]): Sequential of kernel sizes of SPP |
| layers. Defaults to (5, 9, 13). |
| channel_attention (bool): Whether to add channel attention in each |
| stage. Defaults to True. |
| conv_cfg (:obj:`ConfigDict` or dict, optional): Config dict for |
| convolution layer. Defaults to None. |
| norm_cfg (:obj:`ConfigDict` or dict): Dictionary to construct and |
| config norm layer. Defaults to dict(type='BN', requires_grad=True). |
| act_cfg (:obj:`ConfigDict` or dict): Config dict for activation layer. |
| Defaults to dict(type='SiLU'). |
| norm_eval (bool): Whether to set norm layers to eval mode, namely, |
| freeze running stats (mean and var). Note: Effect on Batch Norm |
| and its variants only. |
| init_cfg (:obj:`ConfigDict` or dict or list[dict] or |
| list[:obj:`ConfigDict`]): Initialization config dict. |
| """ |
| |
| |
| arch_settings = { |
| 'P5': [[64, 128, 3, True, False], [128, 256, 6, True, False], |
| [256, 512, 6, True, False], [512, 1024, 3, False, True]], |
| 'P6': [[64, 128, 3, True, False], [128, 256, 6, True, False], |
| [256, 512, 6, True, False], [512, 768, 3, True, False], |
| [768, 1024, 3, False, True]] |
| } |
|
|
| def __init__( |
| self, |
| arch: str = 'P5', |
| deepen_factor: float = 1.0, |
| widen_factor: float = 1.0, |
| out_indices: Sequence[int] = (2, 3, 4), |
| frozen_stages: int = -1, |
| use_depthwise: bool = False, |
| expand_ratio: float = 0.5, |
| arch_ovewrite: dict = None, |
| spp_kernel_sizes: Sequence[int] = (5, 9, 13), |
| channel_attention: bool = True, |
| conv_cfg: OptConfigType = None, |
| norm_cfg: ConfigType = dict(type='BN', momentum=0.03, eps=0.001), |
| act_cfg: ConfigType = dict(type='SiLU'), |
| norm_eval: bool = False, |
| init_cfg: OptMultiConfig = dict( |
| type='Kaiming', |
| layer='Conv2d', |
| a=math.sqrt(5), |
| distribution='uniform', |
| mode='fan_in', |
| nonlinearity='leaky_relu') |
| ) -> None: |
| super().__init__(init_cfg=init_cfg) |
| arch_setting = self.arch_settings[arch] |
| if arch_ovewrite: |
| arch_setting = arch_ovewrite |
| assert set(out_indices).issubset( |
| i for i in range(len(arch_setting) + 1)) |
| if frozen_stages not in range(-1, len(arch_setting) + 1): |
| raise ValueError('frozen_stages must be in range(-1, ' |
| 'len(arch_setting) + 1). But received ' |
| f'{frozen_stages}') |
|
|
| self.out_indices = out_indices |
| self.frozen_stages = frozen_stages |
| self.use_depthwise = use_depthwise |
| self.norm_eval = norm_eval |
| conv = DepthwiseSeparableConvModule if use_depthwise else ConvModule |
| self.stem = nn.Sequential( |
| ConvModule( |
| 3, |
| int(arch_setting[0][0] * widen_factor // 2), |
| 3, |
| padding=1, |
| stride=2, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg), |
| ConvModule( |
| int(arch_setting[0][0] * widen_factor // 2), |
| int(arch_setting[0][0] * widen_factor // 2), |
| 3, |
| padding=1, |
| stride=1, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg), |
| ConvModule( |
| int(arch_setting[0][0] * widen_factor // 2), |
| int(arch_setting[0][0] * widen_factor), |
| 3, |
| padding=1, |
| stride=1, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg)) |
| self.layers = ['stem'] |
|
|
| for i, (in_channels, out_channels, num_blocks, add_identity, |
| use_spp) in enumerate(arch_setting): |
| in_channels = int(in_channels * widen_factor) |
| out_channels = int(out_channels * widen_factor) |
| num_blocks = max(round(num_blocks * deepen_factor), 1) |
| stage = [] |
| conv_layer = conv( |
| in_channels, |
| out_channels, |
| 3, |
| stride=2, |
| padding=1, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg) |
| stage.append(conv_layer) |
| if use_spp: |
| spp = SPPBottleneck( |
| out_channels, |
| out_channels, |
| kernel_sizes=spp_kernel_sizes, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg) |
| stage.append(spp) |
| csp_layer = CSPLayer( |
| out_channels, |
| out_channels, |
| num_blocks=num_blocks, |
| add_identity=add_identity, |
| use_depthwise=use_depthwise, |
| use_cspnext_block=True, |
| expand_ratio=expand_ratio, |
| channel_attention=channel_attention, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| act_cfg=act_cfg) |
| stage.append(csp_layer) |
| self.add_module(f'stage{i + 1}', nn.Sequential(*stage)) |
| self.layers.append(f'stage{i + 1}') |
|
|
| def _freeze_stages(self) -> None: |
| if self.frozen_stages >= 0: |
| for i in range(self.frozen_stages + 1): |
| m = getattr(self, self.layers[i]) |
| m.eval() |
| for param in m.parameters(): |
| param.requires_grad = False |
|
|
| def train(self, mode=True) -> None: |
| super().train(mode) |
| self._freeze_stages() |
| if mode and self.norm_eval: |
| for m in self.modules(): |
| if isinstance(m, _BatchNorm): |
| m.eval() |
|
|
| def forward(self, x: Tuple[Tensor, ...]) -> Tuple[Tensor, ...]: |
| outs = [] |
| for i, layer_name in enumerate(self.layers): |
| layer = getattr(self, layer_name) |
| x = layer(x) |
| if i in self.out_indices: |
| outs.append(x) |
| return tuple(outs) |
|
|
