Implemented config and test file
parent
64fd9187b9
commit
6d2bc718f2
configs/_base_/models
mmpretrain/models/backbones
tests/test_models/test_backbones
|
@ -0,0 +1,12 @@
|
|||
# model settings
|
||||
model = dict(
|
||||
type='ImageClassifier',
|
||||
backbone=dict(type='MobileNetV1', width_mult=1.0),
|
||||
neck=dict(type='GlobalAveragePooling'),
|
||||
head=dict(
|
||||
type='LinearClsHead',
|
||||
num_classes=1000,
|
||||
in_channels=1024,
|
||||
loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
|
||||
topk=(1, 5),
|
||||
))
|
|
@ -14,24 +14,40 @@ class MobileNetV1(BaseBackbone):
|
|||
Args:
|
||||
input_channels (int): The input channels of the image tensor.
|
||||
conv_cfg (dict): Config dict for convolution layer. Default: None.
|
||||
frozen_stages (int): Stages to be frozen (all param fixed). -1 means not freezing any parameters.
|
||||
Default: -1.
|
||||
norm_cfg (dict): Config dict for normalization layer. Default: dict(type='BN').
|
||||
act_cfg (dict): Config dict for activation layer. Default: dict(type='ReLU').
|
||||
norm_eval (bool): Whether to set the normalization layer to evaluation mode. Default: False.
|
||||
with_cp (bool): Use checkpoint or not. Using checkpoint will save some memory while slowing down the
|
||||
training speed. Default: False.
|
||||
frozen_stages (int): Stages to be frozen (all param fixed).
|
||||
-1 means not freezing any parameters.
|
||||
Default: -1.
|
||||
norm_cfg (dict): Config dict for normalization layer.
|
||||
Default: dict(type='BN').
|
||||
act_cfg (dict): Config dict for activation layer.
|
||||
Default: dict(type='ReLU').
|
||||
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 (list[dict]): Initialization config dict. Default: [
|
||||
dict(type='Kaiming', layer=['Conv2d']),
|
||||
dict(type='Constant', val=1, layer=['_BatchNorm', 'GroupNorm'])
|
||||
].
|
||||
"""
|
||||
|
||||
def __init__(self, input_channels, conv_cfg=None, frozen_stages=-1, norm_cfg=dict(type='BN'),
|
||||
act_cfg=dict(type='ReLU'), norm_eval=False, with_cp=False,
|
||||
def __init__(self,
|
||||
input_channels,
|
||||
conv_cfg=None,
|
||||
frozen_stages=-1,
|
||||
norm_cfg=dict(type='BN'),
|
||||
act_cfg=dict(type='ReLU'),
|
||||
norm_eval=False,
|
||||
with_cp=False,
|
||||
init_cfg=[
|
||||
dict(type='Kaiming', layer=['Conv2d']),
|
||||
dict(type='Constant', val=1, layer=['_BatchNorm', 'GroupNorm'])
|
||||
dict(
|
||||
type='Constant',
|
||||
val=1,
|
||||
layer=['_BatchNorm', 'GroupNorm'])
|
||||
]):
|
||||
super(MobileNetV1, self).__init__(init_cfg)
|
||||
self.arch_settings = [[32, 64, 1], [64, 128, 2], [128, 128, 1],
|
||||
|
@ -39,6 +55,9 @@ class MobileNetV1(BaseBackbone):
|
|||
[512, 512, 1], [512, 512, 1], [512, 512, 1],
|
||||
[512, 512, 1], [512, 512, 1], [512, 1024, 2],
|
||||
[1024, 1024, 1]]
|
||||
if frozen_stages not in range(-1, 8):
|
||||
raise ValueError('frozen_stages must be in range(-1, 8). '
|
||||
f'But received {frozen_stages}')
|
||||
self.in_channels = input_channels
|
||||
self.frozen_stages = frozen_stages
|
||||
self.conv_cfg = conv_cfg
|
||||
|
@ -60,19 +79,31 @@ class MobileNetV1(BaseBackbone):
|
|||
act_cfg=self.act_cfg)
|
||||
self.layers.append(layer)
|
||||
|
||||
# Add the rest of the convolution layers to layers according to self.arch_settings
|
||||
for layer_cfg in (self.arch_settings):
|
||||
in_ch, out_ch, stride = layer_cfg
|
||||
self.layers.append(
|
||||
ConvModule(
|
||||
in_channels=in_ch,
|
||||
out_channels=out_ch,
|
||||
kernel_size=3,
|
||||
stride=stride,
|
||||
padding=1,
|
||||
conv_cfg=self.conv_cfg,
|
||||
norm_cfg=self.norm_cfg,
|
||||
act_cfg=self.act_cfg))
|
||||
intermediate_layer = []
|
||||
depthwise_layer = ConvModule(
|
||||
in_channels=in_ch,
|
||||
out_channels=in_ch,
|
||||
kernel_size=3,
|
||||
stride=stride,
|
||||
padding=1,
|
||||
groups=in_ch,
|
||||
conv_cfg=self.conv_cfg,
|
||||
norm_cfg=self.norm_cfg,
|
||||
act_cfg=self.act_cfg)
|
||||
pointwise_layer = ConvModule(
|
||||
in_channels=in_ch,
|
||||
out_channels=out_ch,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=0,
|
||||
conv_cfg=self.conv_cfg,
|
||||
norm_cfg=self.norm_cfg,
|
||||
act_cfg=self.act_cfg)
|
||||
intermediate_layer = nn.Sequential(depthwise_layer,
|
||||
pointwise_layer)
|
||||
self.layers.append(intermediate_layer)
|
||||
self.model = nn.Sequential(*self.layers)
|
||||
|
||||
def forward(self, x):
|
||||
|
|
|
@ -0,0 +1,115 @@
|
|||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import pytest
|
||||
import torch
|
||||
from torch.nn.modules import GroupNorm
|
||||
from torch.nn.modules.batchnorm import _BatchNorm
|
||||
|
||||
from mmpretrain.models.backbones import MobileNetV1
|
||||
|
||||
def is_norm(modules):
|
||||
"""Check if is one of the norms."""
|
||||
if isinstance(modules, (GroupNorm, _BatchNorm)):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def check_norm_state(modules, train_state):
|
||||
"""Check if norm layer is in correct train state."""
|
||||
for mod in modules:
|
||||
if isinstance(mod, _BatchNorm):
|
||||
if mod.training != train_state:
|
||||
return False
|
||||
return True
|
||||
|
||||
def test_mobilenetv1_backbone():
|
||||
with pytest.raises(TypeError):
|
||||
# pretrained must be a string path
|
||||
model = MobileNetV1()
|
||||
model.init_weights(pretrained=0)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
# frozen_stages must in range(-1, 8)
|
||||
MobileNetV1(frozen_stages=8)
|
||||
|
||||
|
||||
# Test MobileNetV2 with first stage frozen
|
||||
frozen_stages = 1
|
||||
model = MobileNetV1(frozen_stages=frozen_stages)
|
||||
model.init_weights()
|
||||
model.train()
|
||||
|
||||
for mod in model.modules():
|
||||
for param in mod.parameters():
|
||||
assert param.requires_grad is False
|
||||
for i in range(1, frozen_stages + 1):
|
||||
layer = getattr(model, f'layer{i}')
|
||||
for mod in layer.modules():
|
||||
if isinstance(mod, _BatchNorm):
|
||||
assert mod.training is False
|
||||
for param in layer.parameters():
|
||||
assert param.requires_grad is False
|
||||
|
||||
# Test MobileNetV2 with norm_eval=True
|
||||
model = MobileNetV1(norm_eval=True)
|
||||
model.init_weights()
|
||||
model.train()
|
||||
|
||||
assert check_norm_state(model.modules(), False)
|
||||
|
||||
|
||||
# Test MobileNetV2 forward with dict(type='ReLU')
|
||||
model = MobileNetV1(act_cfg=dict(type='ReLU'))
|
||||
model.init_weights()
|
||||
model.train()
|
||||
|
||||
imgs = torch.randn(1, 3, 224, 224)
|
||||
feat = model(imgs)
|
||||
assert len(feat) == 7
|
||||
assert feat[0].shape == torch.Size((1, 16, 112, 112))
|
||||
assert feat[1].shape == torch.Size((1, 24, 56, 56))
|
||||
assert feat[2].shape == torch.Size((1, 32, 28, 28))
|
||||
assert feat[3].shape == torch.Size((1, 64, 14, 14))
|
||||
assert feat[4].shape == torch.Size((1, 96, 14, 14))
|
||||
assert feat[5].shape == torch.Size((1, 160, 7, 7))
|
||||
assert feat[6].shape == torch.Size((1, 320, 7, 7))
|
||||
|
||||
# Test MobileNetV2 with BatchNorm forward
|
||||
model = MobileNetV1()
|
||||
for m in model.modules():
|
||||
if is_norm(m):
|
||||
assert isinstance(m, _BatchNorm)
|
||||
model.init_weights()
|
||||
model.train()
|
||||
|
||||
imgs = torch.randn(1, 3, 224, 224)
|
||||
feat = model(imgs)
|
||||
assert len(feat) == 7
|
||||
assert feat[0].shape == torch.Size((1, 16, 112, 112))
|
||||
assert feat[1].shape == torch.Size((1, 24, 56, 56))
|
||||
assert feat[2].shape == torch.Size((1, 32, 28, 28))
|
||||
assert feat[3].shape == torch.Size((1, 64, 14, 14))
|
||||
assert feat[4].shape == torch.Size((1, 96, 14, 14))
|
||||
assert feat[5].shape == torch.Size((1, 160, 7, 7))
|
||||
assert feat[6].shape == torch.Size((1, 320, 7, 7))
|
||||
|
||||
# Test MobileNetV2 with GroupNorm forward
|
||||
model = MobileNetV1(
|
||||
norm_cfg=dict(type='GN', num_groups=2, requires_grad=True))
|
||||
for m in model.modules():
|
||||
if is_norm(m):
|
||||
assert isinstance(m, GroupNorm)
|
||||
model.init_weights()
|
||||
model.train()
|
||||
|
||||
imgs = torch.randn(1, 3, 224, 224)
|
||||
feat = model(imgs)
|
||||
assert len(feat) == 7
|
||||
assert feat[0].shape == torch.Size((1, 16, 112, 112))
|
||||
assert feat[1].shape == torch.Size((1, 24, 56, 56))
|
||||
assert feat[2].shape == torch.Size((1, 32, 28, 28))
|
||||
assert feat[3].shape == torch.Size((1, 64, 14, 14))
|
||||
assert feat[4].shape == torch.Size((1, 96, 14, 14))
|
||||
assert feat[5].shape == torch.Size((1, 160, 7, 7))
|
||||
assert feat[6].shape == torch.Size((1, 320, 7, 7))
|
||||
|
||||
|
Loading…
Reference in New Issue