mmsegmentation框架下给mobilenetv2主干加入SKNet
时间: 2023-10-01 15:07:54 浏览: 54
在`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`结构。