Shortcuts

Source code for mmpose.models.backbones.vipnas_mbv3

# 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


[docs]@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
[docs] 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
[docs] 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()
Read the Docs v: latest
Versions
latest
0.x
dev-1.x
Downloads
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.