Source code for mmpose.models.backbones.scnet
# Copyright (c) OpenMMLab. All rights reserved.
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 mmengine.model import BaseModule
from mmpose.registry import MODELS
from .resnet import Bottleneck, ResNet
class SCConv(BaseModule):
"""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')
init_cfg (dict or list[dict], optional): Initialization config dict.
Default: None
"""
def __init__(self,
in_channels,
out_channels,
stride,
pooling_r,
conv_cfg=None,
norm_cfg=dict(type='BN', momentum=0.1),
init_cfg=None):
# Protect mutable default arguments
norm_cfg = copy.deepcopy(norm_cfg)
super().__init__(init_cfg=init_cfg)
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
[docs]@MODELS.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)