|
|
|
import copy |
|
|
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
import torch.utils.checkpoint as cp |
|
from mmcv.cnn import build_conv_layer, build_norm_layer |
|
|
|
from ..builder import BACKBONES |
|
from .resnet import Bottleneck, ResNet |
|
|
|
|
|
class SCConv(nn.Module): |
|
"""SCConv (Self-calibrated Convolution) |
|
|
|
Args: |
|
in_channels (int): The input channels of the SCConv. |
|
out_channels (int): The output channel of the SCConv. |
|
stride (int): stride of SCConv. |
|
pooling_r (int): size of pooling for scconv. |
|
conv_cfg (dict): dictionary to construct and config conv layer. |
|
Default: None |
|
norm_cfg (dict): dictionary to construct and config norm layer. |
|
Default: dict(type='BN') |
|
""" |
|
|
|
def __init__(self, |
|
in_channels, |
|
out_channels, |
|
stride, |
|
pooling_r, |
|
conv_cfg=None, |
|
norm_cfg=dict(type='BN', momentum=0.1)): |
|
|
|
norm_cfg = copy.deepcopy(norm_cfg) |
|
super().__init__() |
|
|
|
assert in_channels == out_channels |
|
|
|
self.k2 = nn.Sequential( |
|
nn.AvgPool2d(kernel_size=pooling_r, stride=pooling_r), |
|
build_conv_layer( |
|
conv_cfg, |
|
in_channels, |
|
in_channels, |
|
kernel_size=3, |
|
stride=1, |
|
padding=1, |
|
bias=False), |
|
build_norm_layer(norm_cfg, in_channels)[1], |
|
) |
|
self.k3 = nn.Sequential( |
|
build_conv_layer( |
|
conv_cfg, |
|
in_channels, |
|
in_channels, |
|
kernel_size=3, |
|
stride=1, |
|
padding=1, |
|
bias=False), |
|
build_norm_layer(norm_cfg, in_channels)[1], |
|
) |
|
self.k4 = nn.Sequential( |
|
build_conv_layer( |
|
conv_cfg, |
|
in_channels, |
|
in_channels, |
|
kernel_size=3, |
|
stride=stride, |
|
padding=1, |
|
bias=False), |
|
build_norm_layer(norm_cfg, out_channels)[1], |
|
nn.ReLU(inplace=True), |
|
) |
|
|
|
def forward(self, x): |
|
"""Forward function.""" |
|
identity = x |
|
|
|
out = torch.sigmoid( |
|
torch.add(identity, F.interpolate(self.k2(x), |
|
identity.size()[2:]))) |
|
out = torch.mul(self.k3(x), out) |
|
out = self.k4(out) |
|
|
|
return out |
|
|
|
|
|
class SCBottleneck(Bottleneck): |
|
"""SC(Self-calibrated) Bottleneck. |
|
|
|
Args: |
|
in_channels (int): The input channels of the SCBottleneck block. |
|
out_channels (int): The output channel of the SCBottleneck block. |
|
""" |
|
|
|
pooling_r = 4 |
|
|
|
def __init__(self, in_channels, out_channels, **kwargs): |
|
super().__init__(in_channels, out_channels, **kwargs) |
|
self.mid_channels = out_channels // self.expansion // 2 |
|
|
|
self.norm1_name, norm1 = build_norm_layer( |
|
self.norm_cfg, self.mid_channels, postfix=1) |
|
self.norm2_name, norm2 = build_norm_layer( |
|
self.norm_cfg, self.mid_channels, postfix=2) |
|
self.norm3_name, norm3 = build_norm_layer( |
|
self.norm_cfg, out_channels, postfix=3) |
|
|
|
self.conv1 = build_conv_layer( |
|
self.conv_cfg, |
|
in_channels, |
|
self.mid_channels, |
|
kernel_size=1, |
|
stride=1, |
|
bias=False) |
|
self.add_module(self.norm1_name, norm1) |
|
|
|
self.k1 = nn.Sequential( |
|
build_conv_layer( |
|
self.conv_cfg, |
|
self.mid_channels, |
|
self.mid_channels, |
|
kernel_size=3, |
|
stride=self.stride, |
|
padding=1, |
|
bias=False), |
|
build_norm_layer(self.norm_cfg, self.mid_channels)[1], |
|
nn.ReLU(inplace=True)) |
|
|
|
self.conv2 = build_conv_layer( |
|
self.conv_cfg, |
|
in_channels, |
|
self.mid_channels, |
|
kernel_size=1, |
|
stride=1, |
|
bias=False) |
|
self.add_module(self.norm2_name, norm2) |
|
|
|
self.scconv = SCConv(self.mid_channels, self.mid_channels, self.stride, |
|
self.pooling_r, self.conv_cfg, self.norm_cfg) |
|
|
|
self.conv3 = build_conv_layer( |
|
self.conv_cfg, |
|
self.mid_channels * 2, |
|
out_channels, |
|
kernel_size=1, |
|
stride=1, |
|
bias=False) |
|
self.add_module(self.norm3_name, norm3) |
|
|
|
def forward(self, x): |
|
"""Forward function.""" |
|
|
|
def _inner_forward(x): |
|
identity = x |
|
|
|
out_a = self.conv1(x) |
|
out_a = self.norm1(out_a) |
|
out_a = self.relu(out_a) |
|
|
|
out_a = self.k1(out_a) |
|
|
|
out_b = self.conv2(x) |
|
out_b = self.norm2(out_b) |
|
out_b = self.relu(out_b) |
|
|
|
out_b = self.scconv(out_b) |
|
|
|
out = self.conv3(torch.cat([out_a, out_b], dim=1)) |
|
out = self.norm3(out) |
|
|
|
if self.downsample is not None: |
|
identity = self.downsample(x) |
|
|
|
out += identity |
|
|
|
return out |
|
|
|
if self.with_cp and x.requires_grad: |
|
out = cp.checkpoint(_inner_forward, x) |
|
else: |
|
out = _inner_forward(x) |
|
|
|
out = self.relu(out) |
|
|
|
return out |
|
|
|
|
|
@BACKBONES.register_module() |
|
class SCNet(ResNet): |
|
"""SCNet backbone. |
|
|
|
Improving Convolutional Networks with Self-Calibrated Convolutions, |
|
Jiang-Jiang Liu, Qibin Hou, Ming-Ming Cheng, Changhu Wang, Jiashi Feng, |
|
IEEE CVPR, 2020. |
|
http://mftp.mmcheng.net/Papers/20cvprSCNet.pdf |
|
|
|
Args: |
|
depth (int): Depth of scnet, from {50, 101}. |
|
in_channels (int): Number of input image channels. Normally 3. |
|
base_channels (int): Number of base channels of hidden layer. |
|
num_stages (int): SCNet stages, normally 4. |
|
strides (Sequence[int]): Strides of the first block of each stage. |
|
dilations (Sequence[int]): Dilation of each stage. |
|
out_indices (Sequence[int]): Output from which stages. |
|
style (str): `pytorch` or `caffe`. If set to "pytorch", the stride-two |
|
layer is the 3x3 conv layer, otherwise the stride-two layer is |
|
the first 1x1 conv layer. |
|
deep_stem (bool): Replace 7x7 conv in input stem with 3 3x3 conv |
|
avg_down (bool): Use AvgPool instead of stride conv when |
|
downsampling in the bottleneck. |
|
frozen_stages (int): Stages to be frozen (stop grad and set eval mode). |
|
-1 means not freezing any parameters. |
|
norm_cfg (dict): Dictionary to construct and config norm layer. |
|
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. |
|
with_cp (bool): Use checkpoint or not. Using checkpoint will save some |
|
memory while slowing down the training speed. |
|
zero_init_residual (bool): Whether to use zero init for last norm layer |
|
in resblocks to let them behave as identity. |
|
|
|
Example: |
|
>>> from mmpose.models import SCNet |
|
>>> import torch |
|
>>> self = SCNet(depth=50, out_indices=(0, 1, 2, 3)) |
|
>>> self.eval() |
|
>>> inputs = torch.rand(1, 3, 224, 224) |
|
>>> level_outputs = self.forward(inputs) |
|
>>> for level_out in level_outputs: |
|
... print(tuple(level_out.shape)) |
|
(1, 256, 56, 56) |
|
(1, 512, 28, 28) |
|
(1, 1024, 14, 14) |
|
(1, 2048, 7, 7) |
|
""" |
|
|
|
arch_settings = { |
|
50: (SCBottleneck, [3, 4, 6, 3]), |
|
101: (SCBottleneck, [3, 4, 23, 3]) |
|
} |
|
|
|
def __init__(self, depth, **kwargs): |
|
if depth not in self.arch_settings: |
|
raise KeyError(f'invalid depth {depth} for SCNet') |
|
super().__init__(depth, **kwargs) |
|
|
