| |
| from typing import Optional |
|
|
| from mmcv.cnn import build_conv_layer, build_norm_layer |
| from mmengine.model import BaseModule, Sequential |
| from torch import Tensor |
| from torch import nn as nn |
|
|
| from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig |
|
|
|
|
| class ResLayer(Sequential): |
| """ResLayer to build ResNet style backbone. |
| |
| Args: |
| block (nn.Module): block used to build ResLayer. |
| inplanes (int): inplanes of block. |
| planes (int): planes of block. |
| num_blocks (int): number of blocks. |
| stride (int): stride of the first block. Defaults to 1 |
| avg_down (bool): Use AvgPool instead of stride conv when |
| downsampling in the bottleneck. Defaults to False |
| conv_cfg (dict): dictionary to construct and config conv layer. |
| Defaults to None |
| norm_cfg (dict): dictionary to construct and config norm layer. |
| Defaults to dict(type='BN') |
| downsample_first (bool): Downsample at the first block or last block. |
| False for Hourglass, True for ResNet. Defaults to True |
| """ |
|
|
| def __init__(self, |
| block: BaseModule, |
| inplanes: int, |
| planes: int, |
| num_blocks: int, |
| stride: int = 1, |
| avg_down: bool = False, |
| conv_cfg: OptConfigType = None, |
| norm_cfg: ConfigType = dict(type='BN'), |
| downsample_first: bool = True, |
| **kwargs) -> None: |
| self.block = block |
|
|
| downsample = None |
| if stride != 1 or inplanes != planes * block.expansion: |
| downsample = [] |
| conv_stride = stride |
| if avg_down: |
| conv_stride = 1 |
| downsample.append( |
| nn.AvgPool2d( |
| kernel_size=stride, |
| stride=stride, |
| ceil_mode=True, |
| count_include_pad=False)) |
| downsample.extend([ |
| build_conv_layer( |
| conv_cfg, |
| inplanes, |
| planes * block.expansion, |
| kernel_size=1, |
| stride=conv_stride, |
| bias=False), |
| build_norm_layer(norm_cfg, planes * block.expansion)[1] |
| ]) |
| downsample = nn.Sequential(*downsample) |
|
|
| layers = [] |
| if downsample_first: |
| layers.append( |
| block( |
| inplanes=inplanes, |
| planes=planes, |
| stride=stride, |
| downsample=downsample, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| **kwargs)) |
| inplanes = planes * block.expansion |
| for _ in range(1, num_blocks): |
| layers.append( |
| block( |
| inplanes=inplanes, |
| planes=planes, |
| stride=1, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| **kwargs)) |
|
|
| else: |
| for _ in range(num_blocks - 1): |
| layers.append( |
| block( |
| inplanes=inplanes, |
| planes=inplanes, |
| stride=1, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| **kwargs)) |
| layers.append( |
| block( |
| inplanes=inplanes, |
| planes=planes, |
| stride=stride, |
| downsample=downsample, |
| conv_cfg=conv_cfg, |
| norm_cfg=norm_cfg, |
| **kwargs)) |
| super().__init__(*layers) |
|
|
|
|
| class SimplifiedBasicBlock(BaseModule): |
| """Simplified version of original basic residual block. This is used in |
| `SCNet <https://arxiv.org/abs/2012.10150>`_. |
| |
| - Norm layer is now optional |
| - Last ReLU in forward function is removed |
| """ |
| expansion = 1 |
|
|
| def __init__(self, |
| inplanes: int, |
| planes: int, |
| stride: int = 1, |
| dilation: int = 1, |
| downsample: Optional[Sequential] = None, |
| style: ConfigType = 'pytorch', |
| with_cp: bool = False, |
| conv_cfg: OptConfigType = None, |
| norm_cfg: ConfigType = dict(type='BN'), |
| dcn: OptConfigType = None, |
| plugins: OptConfigType = None, |
| init_cfg: OptMultiConfig = None) -> None: |
| super().__init__(init_cfg=init_cfg) |
| assert dcn is None, 'Not implemented yet.' |
| assert plugins is None, 'Not implemented yet.' |
| assert not with_cp, 'Not implemented yet.' |
| self.with_norm = norm_cfg is not None |
| with_bias = True if norm_cfg is None else False |
| self.conv1 = build_conv_layer( |
| conv_cfg, |
| inplanes, |
| planes, |
| 3, |
| stride=stride, |
| padding=dilation, |
| dilation=dilation, |
| bias=with_bias) |
| if self.with_norm: |
| self.norm1_name, norm1 = build_norm_layer( |
| norm_cfg, planes, postfix=1) |
| self.add_module(self.norm1_name, norm1) |
| self.conv2 = build_conv_layer( |
| conv_cfg, planes, planes, 3, padding=1, bias=with_bias) |
| if self.with_norm: |
| self.norm2_name, norm2 = build_norm_layer( |
| norm_cfg, planes, postfix=2) |
| self.add_module(self.norm2_name, norm2) |
|
|
| self.relu = nn.ReLU(inplace=True) |
| self.downsample = downsample |
| self.stride = stride |
| self.dilation = dilation |
| self.with_cp = with_cp |
|
|
| @property |
| def norm1(self) -> Optional[BaseModule]: |
| """nn.Module: normalization layer after the first convolution layer""" |
| return getattr(self, self.norm1_name) if self.with_norm else None |
|
|
| @property |
| def norm2(self) -> Optional[BaseModule]: |
| """nn.Module: normalization layer after the second convolution layer""" |
| return getattr(self, self.norm2_name) if self.with_norm else None |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| """Forward function for SimplifiedBasicBlock.""" |
|
|
| identity = x |
|
|
| out = self.conv1(x) |
| if self.with_norm: |
| out = self.norm1(out) |
| out = self.relu(out) |
|
|
| out = self.conv2(out) |
| if self.with_norm: |
| out = self.norm2(out) |
|
|
| if self.downsample is not None: |
| identity = self.downsample(x) |
|
|
| out += identity |
|
|
| return out |
|
|
