mmpose/models/backbones/vipnas_mbv3.py

# Copyright (c) OpenMMLab. All rights reserved.
import copy

from mmcv.cnn import ConvModule
from torch.nn.modules.batchnorm import _BatchNorm

from mmpose.registry import MODELS
from .base_backbone import BaseBackbone
from .utils import InvertedResidual


@MODELS.register_module()
class ViPNAS_MobileNetV3(BaseBackbone):
    """ViPNAS_MobileNetV3 backbone.

    "ViPNAS: Efficient Video Pose Estimation via Neural Architecture Search"
    More details can be found in the `paper
    <https://arxiv.org/abs/2105.10154>`__ .

    Args:
        wid (list(int)): Searched width config for each stage.
        expan (list(int)): Searched expansion ratio config for each stage.
        dep (list(int)): Searched depth config for each stage.
        ks (list(int)): Searched kernel size config for each stage.
        group (list(int)): Searched group number config for each stage.
        att (list(bool)): Searched attention config for each stage.
        stride (list(int)): Stride config for each stage.
        act (list(dict)): Activation config for each stage.
        conv_cfg (dict): Config dict for convolution layer.
            Default: None, which means using conv2d.
        norm_cfg (dict): Config dict for normalization layer.
            Default: dict(type='BN').
        frozen_stages (int): Stages to be frozen (all param fixed).
            Default: -1, which means not freezing any parameters.
        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. Default: False.
        with_cp (bool): Use checkpoint or not. Using checkpoint will save
            some memory while slowing down the training speed.
            Default: False.
        init_cfg (dict or list[dict], optional): Initialization config dict.
            Default:
            ``[
                dict(type='Normal', std=0.001, layer=['Conv2d']),
                dict(
                    type='Constant',
                    val=1,
                    layer=['_BatchNorm', 'GroupNorm'])
            ]``
    """

    def __init__(
        self,
        wid=[16, 16, 24, 40, 80, 112, 160],
        expan=[None, 1, 5, 4, 5, 5, 6],
        dep=[None, 1, 4, 4, 4, 4, 4],
        ks=[3, 3, 7, 7, 5, 7, 5],
        group=[None, 8, 120, 20, 100, 280, 240],
        att=[None, True, True, False, True, True, True],
        stride=[2, 1, 2, 2, 2, 1, 2],
        act=['HSwish', 'ReLU', 'ReLU', 'ReLU', 'HSwish', 'HSwish', 'HSwish'],
        conv_cfg=None,
        norm_cfg=dict(type='BN'),
        frozen_stages=-1,
        norm_eval=False,
        with_cp=False,
        init_cfg=[
            dict(type='Normal', std=0.001, layer=['Conv2d']),
            dict(type='Constant', val=1, layer=['_BatchNorm', 'GroupNorm'])
        ],
    ):
        # Protect mutable default arguments
        norm_cfg = copy.deepcopy(norm_cfg)
        super().__init__(init_cfg=init_cfg)
        self.wid = wid
        self.expan = expan
        self.dep = dep
        self.ks = ks
        self.group = group
        self.att = att
        self.stride = stride
        self.act = act
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.frozen_stages = frozen_stages
        self.norm_eval = norm_eval
        self.with_cp = with_cp

        self.conv1 = ConvModule(
            in_channels=3,
            out_channels=self.wid[0],
            kernel_size=self.ks[0],
            stride=self.stride[0],
            padding=self.ks[0] // 2,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=dict(type=self.act[0]))

        self.layers = self._make_layer()

    def _make_layer(self):
        layers = []
        layer_index = 0
        for i, dep in enumerate(self.dep[1:]):
            mid_channels = self.wid[i + 1] * self.expan[i + 1]

            if self.att[i + 1]:
                se_cfg = dict(
                    channels=mid_channels,
                    ratio=4,
                    act_cfg=(dict(type='ReLU'),
                             dict(type='HSigmoid', bias=1.0, divisor=2.0)))
            else:
                se_cfg = None

            if self.expan[i + 1] == 1:
                with_expand_conv = False
            else:
                with_expand_conv = True

            for j in range(dep):
                if j == 0:
                    stride = self.stride[i + 1]
                    in_channels = self.wid[i]
                else:
                    stride = 1
                    in_channels = self.wid[i + 1]

                layer = InvertedResidual(
                    in_channels=in_channels,
                    out_channels=self.wid[i + 1],
                    mid_channels=mid_channels,
                    kernel_size=self.ks[i + 1],
                    groups=self.group[i + 1],
                    stride=stride,
                    se_cfg=se_cfg,
                    with_expand_conv=with_expand_conv,
                    conv_cfg=self.conv_cfg,
                    norm_cfg=self.norm_cfg,
                    act_cfg=dict(type=self.act[i + 1]),
                    with_cp=self.with_cp)
                layer_index += 1
                layer_name = f'layer{layer_index}'
                self.add_module(layer_name, layer)
                layers.append(layer_name)
        return layers

    def forward(self, x):
        x = self.conv1(x)

        for i, layer_name in enumerate(self.layers):
            layer = getattr(self, layer_name)
            x = layer(x)

        return (x, )

    def _freeze_stages(self):
        if self.frozen_stages >= 0:
            for param in self.conv1.parameters():
                param.requires_grad = False
        for i in range(1, self.frozen_stages + 1):
            layer = getattr(self, f'layer{i}')
            layer.eval()
            for param in layer.parameters():
                param.requires_grad = False

    def train(self, mode=True):
        super().train(mode)
        self._freeze_stages()
        if mode and self.norm_eval:
            for m in self.modules():
                if isinstance(m, _BatchNorm):
                    m.eval()