mmsegmentation框架下给mobilenetv2主干加入SKNet

在`mmsegmentation`框架中，可以通过修改`mmseg/models/backbones/mobilenet_v2.py`文件来给`MobileNetV2`主干加入`SKNet`。下面是实现的步骤： 1. 首先需要导入`SKNet`的相关模块，可以在文件最上面添加以下代码： ```python from mmcv.cnn import build_conv_layer, build_norm_layer, build_plugin_layer from ..builder import BACKBONES ``` 2. 接着在`MobileNetV2`类中定义`SKNet`的结构，可以在`_make_stem_layer`函数中加入以下代码： ```python def _make_stem_layer(self, in_channels, stem_channels): layers = [] layers.append(ConvModule( in_channels, stem_channels, 3, stride=2, padding=1, bias=False, norm_cfg=dict(type='BN', momentum=0.1, eps=1e-5), activation='relu', inplace=True)) in_channels = stem_channels layers.append(ConvModule( in_channels, in_channels, 3, stride=1, padding=1, bias=False, norm_cfg=dict(type='BN', momentum=0.1, eps=1e-5), activation='relu', inplace=True)) # add SKNet module channels = in_channels mid_channels = channels // 2 squeeze_channels = max(1, mid_channels // 8) layers.append( build_plugin_layer(dict( type='SKConv', channels=channels, squeeze_channels=squeeze_channels, kernel_size=3, stride=1, padding=1, dilation=1, groups=32, sk_mode='two', norm_cfg=dict(type='BN', momentum=0.1, eps=1e-5), act_cfg=dict(type='ReLU', inplace=True), ), [build_conv_layer( dict(type='Conv2d'), channels, channels, kernel_size=3, stride=1, padding=1, bias=False), build_norm_layer(dict(type='BN', momentum=0.1, eps=1e-5), channels)[1]])) return nn.Sequential(*layers) ``` 3. 最后在`BACKBONES`中注册`MobileNetV2`主干即可。完整代码如下： ```python import torch.nn as nn from mmcv.cnn import ConvModule from mmcv.cnn import build_conv_layer, build_norm_layer, build_plugin_layer from ..builder import BACKBONES @BACKBONES.register_module() class MobileNetV2(nn.Module): def __init__(self, widen_factor=1.0, output_stride=32, norm_cfg=dict(type='BN', momentum=0.1, eps=1e-5), with_cp=False, ): super(MobileNetV2, self).__init__() assert output_stride in [8, 16, 32] self.output_stride = output_stride self.with_cp = with_cp self.norm_cfg = norm_cfg input_channel = int(32 * widen_factor) self.stem = self._make_stem_layer(3, input_channel) self.layer1 = self._make_layer( input_channel, int(16 * widen_factor), 1, 1, 16, 2) self.layer2 = self._make_layer( int(16 * widen_factor), int(24 * widen_factor), 2, 6, 16, 2) self.layer3 = self._make_layer( int(24 * widen_factor), int(32 * widen_factor), 3, 6, 24, 2) self.layer4 = self._make_layer( int(32 * widen_factor), int(64 * widen_factor), 4, 6, 32, 2) self.layer5 = self._make_layer( int(64 * widen_factor), int(96 * widen_factor), 3, 6, 64, 1) self.layer6 = self._make_layer( int(96 * widen_factor), int(160 * widen_factor), 3, 6, 96, 1) self.layer7 = self._make_layer( int(160 * widen_factor), int(320 * widen_factor), 1, 6, 160, 1) if self.output_stride == 8: self.layer2[0].conv2.stride = (1, 1) self.layer2[0].downsample[0].stride = (1, 1) self.layer3[0].conv2.stride = (1, 1) self.layer3[0].downsample[0].stride = (1, 1) elif self.output_stride == 16: self.layer3[0].conv2.stride = (1, 1) self.layer3[0].downsample[0].stride = (1, 1) self._freeze_stages() def _make_stem_layer(self, in_channels, stem_channels): layers = [] layers.append(ConvModule( in_channels, stem_channels, 3, stride=2, padding=1, bias=False, norm_cfg=dict(type='BN', momentum=0.1, eps=1e-5), activation='relu', inplace=True)) in_channels = stem_channels layers.append(ConvModule( in_channels, in_channels, 3, stride=1, padding=1, bias=False, norm_cfg=dict(type='BN', momentum=0.1, eps=1e-5), activation='relu', inplace=True)) # add SKNet module channels = in_channels mid_channels = channels // 2 squeeze_channels = max(1, mid_channels // 8) layers.append( build_plugin_layer(dict( type='SKConv', channels=channels, squeeze_channels=squeeze_channels, kernel_size=3, stride=1, padding=1, dilation=1, groups=32, sk_mode='two', norm_cfg=dict(type='BN', momentum=0.1, eps=1e-5), act_cfg=dict(type='ReLU', inplace=True), ), [build_conv_layer( dict(type='Conv2d'), channels, channels, kernel_size=3, stride=1, padding=1, bias=False), build_norm_layer(dict(type='BN', momentum=0.1, eps=1e-5), channels)[1]])) return nn.Sequential(*layers) def _make_layer(self, in_channels, out_channels, num_blocks, stride, t, dilation=1): layers = [] layers.append(InvertedResidual( in_channels, out_channels, stride, expand_ratio=t, dilation=dilation, norm_cfg=self.norm_cfg, with_cp=self.with_cp)) for i in range(1, num_blocks): layers.append(InvertedResidual( out_channels, out_channels, 1, expand_ratio=t, dilation=dilation, norm_cfg=self.norm_cfg, with_cp=self.with_cp)) return nn.Sequential(*layers) def _freeze_stages(self): if self.output_stride == 8: self.stem.eval() for param in self.stem.parameters(): param.requires_grad = False elif self.output_stride == 16: self.stem[0].eval() self.stem[1].eval() for param in self.stem[0].parameters(): param.requires_grad = False for param in self.stem[1].parameters(): param.requires_grad = False ``` 这样就给`MobileNetV2`主干加入了`SKNet`结构。