pull/2/head
lixiaojie 2020-06-04 02:42:53 +08:00
parent b392ba44cc
commit fa49911440
2 changed files with 56 additions and 64 deletions

View File

@ -22,13 +22,12 @@ def conv3x3(in_planes, out_planes, stride=1, dilation=1):
def conv_1x1_bn(inp, oup, activation=nn.ReLU6): def conv_1x1_bn(inp, oup, activation=nn.ReLU6):
return nn.Sequential( return nn.Sequential(
nn.Conv2d(inp, oup, 1, 1, 0, bias=False), nn.Conv2d(inp, oup, 1, 1, 0, bias=False), nn.BatchNorm2d(oup),
nn.BatchNorm2d(oup), activation(inplace=True))
activation(inplace=True)
)
class ConvBNReLU(nn.Sequential): class ConvBNReLU(nn.Sequential):
def __init__(self, def __init__(self,
in_planes, in_planes,
out_planes, out_planes,
@ -39,16 +38,15 @@ class ConvBNReLU(nn.Sequential):
padding = (kernel_size - 1) // 2 padding = (kernel_size - 1) // 2
super(ConvBNReLU, self).__init__( super(ConvBNReLU, self).__init__(
nn.Conv2d(in_planes, nn.Conv2d(
out_planes, in_planes,
kernel_size, out_planes,
stride, kernel_size,
padding, stride,
groups=groups, padding,
bias=False), groups=groups,
nn.BatchNorm2d(out_planes), bias=False), nn.BatchNorm2d(out_planes),
activation(inplace=True) activation(inplace=True))
)
def _make_divisible(v, divisor, min_value=None): def _make_divisible(v, divisor, min_value=None):
@ -62,6 +60,7 @@ def _make_divisible(v, divisor, min_value=None):
class InvertedResidual(nn.Module): class InvertedResidual(nn.Module):
def __init__(self, def __init__(self,
inplanes, inplanes,
outplanes, outplanes,
@ -79,17 +78,18 @@ class InvertedResidual(nn.Module):
layers = [] layers = []
if expand_ratio != 1: if expand_ratio != 1:
# pw # pw
layers.append(ConvBNReLU(inplanes, layers.append(
hidden_dim, ConvBNReLU(
kernel_size=1, inplanes, hidden_dim, kernel_size=1,
activation=activation)) activation=activation))
layers.extend([ layers.extend([
# dw # dw
ConvBNReLU(hidden_dim, ConvBNReLU(
hidden_dim, hidden_dim,
stride=stride, hidden_dim,
groups=hidden_dim, stride=stride,
activation=activation), groups=hidden_dim,
activation=activation),
# pw-linear # pw-linear
nn.Conv2d(hidden_dim, outplanes, 1, 1, 0, bias=False), nn.Conv2d(hidden_dim, outplanes, 1, 1, 0, bias=False),
nn.BatchNorm2d(outplanes), nn.BatchNorm2d(outplanes),
@ -97,6 +97,7 @@ class InvertedResidual(nn.Module):
self.conv = nn.Sequential(*layers) self.conv = nn.Sequential(*layers)
def forward(self, x): def forward(self, x):
def _inner_forward(x): def _inner_forward(x):
if self.use_res_connect: if self.use_res_connect:
return x + self.conv(x) return x + self.conv(x)
@ -122,15 +123,23 @@ def make_inverted_res_layer(block,
layers = [] layers = []
for i in range(num_blocks): for i in range(num_blocks):
if i == 0: if i == 0:
layers.append(block(inplanes, planes, stride, layers.append(
expand_ratio=expand_ratio, block(
activation=activation, inplanes,
with_cp=with_cp)) planes,
stride,
expand_ratio=expand_ratio,
activation=activation,
with_cp=with_cp))
else: else:
layers.append(block(inplanes, planes, 1, layers.append(
expand_ratio=expand_ratio, block(
activation=activation, inplanes,
with_cp=with_cp)) planes,
1,
expand_ratio=expand_ratio,
activation=activation,
with_cp=with_cp))
inplanes = planes inplanes = planes
return nn.Sequential(*layers) return nn.Sequential(*layers)
@ -162,15 +171,10 @@ class MobileNetv2(BaseBackbone):
super(MobileNetv2, self).__init__() super(MobileNetv2, self).__init__()
block = InvertedResidual block = InvertedResidual
# expand_ratio, out_channel, n, stride # expand_ratio, out_channel, n, stride
inverted_residual_setting = [ inverted_residual_setting = [[1, 16, 1, 1], [6, 24, 2,
[1, 16, 1, 1], 2], [6, 32, 3, 2],
[6, 24, 2, 2], [6, 64, 4, 2], [6, 96, 3, 1],
[6, 32, 3, 2], [6, 160, 3, 2], [6, 320, 1, 1]]
[6, 64, 4, 2],
[6, 96, 3, 1],
[6, 160, 3, 2],
[6, 320, 1, 1]
]
self.widen_factor = widen_factor self.widen_factor = widen_factor
if isinstance(activation, str): if isinstance(activation, str):
activation = eval(activation) activation = eval(activation)
@ -211,9 +215,8 @@ class MobileNetv2(BaseBackbone):
self.out_channel = int(self.out_channel * widen_factor) \ self.out_channel = int(self.out_channel * widen_factor) \
if widen_factor > 1.0 else self.out_channel if widen_factor > 1.0 else self.out_channel
self.conv_last = nn.Conv2d(self.inplanes, self.conv_last = nn.Conv2d(
self.out_channel, self.inplanes, self.out_channel, 1, 1, 0, bias=False)
1, 1, 0, bias=False)
self.bn_last = nn.BatchNorm2d(self.out_channel) self.bn_last = nn.BatchNorm2d(self.out_channel)
self.feat_dim = self.out_channel self.feat_dim = self.out_channel

View File

@ -10,7 +10,7 @@ from mmcls.models.backbones.mobilenet_v2 import InvertedResidual
def is_block(modules): def is_block(modules):
"""Check if is ResNet building block.""" """Check if is ResNet building block."""
if isinstance(modules, (InvertedResidual,)): if isinstance(modules, (InvertedResidual, )):
return True return True
return False return False
@ -35,40 +35,30 @@ def test_mobilenetv2_invertedresidual():
with pytest.raises(AssertionError): with pytest.raises(AssertionError):
# stride must be in [1, 2] # stride must be in [1, 2]
InvertedResidual(64, 16, InvertedResidual(64, 16, stride=3, expand_ratio=6)
stride=3, expand_ratio=6)
# Test InvertedResidual with checkpoint forward, stride=1 # Test InvertedResidual with checkpoint forward, stride=1
block = InvertedResidual(64, 16, block = InvertedResidual(64, 16, stride=1, expand_ratio=6)
stride=1,
expand_ratio=6)
x = torch.randn(1, 64, 56, 56) x = torch.randn(1, 64, 56, 56)
x_out = block(x) x_out = block(x)
assert x_out.shape == torch.Size([1, 16, 56, 56]) assert x_out.shape == torch.Size([1, 16, 56, 56])
# Test InvertedResidual with checkpoint forward, stride=2 # Test InvertedResidual with checkpoint forward, stride=2
block = InvertedResidual(64, 16, block = InvertedResidual(64, 16, stride=2, expand_ratio=6)
stride=2,
expand_ratio=6)
x = torch.randn(1, 64, 56, 56) x = torch.randn(1, 64, 56, 56)
x_out = block(x) x_out = block(x)
assert x_out.shape == torch.Size([1, 16, 28, 28]) assert x_out.shape == torch.Size([1, 16, 28, 28])
# Test InvertedResidual with checkpoint forward # Test InvertedResidual with checkpoint forward
block = InvertedResidual(64, 16, block = InvertedResidual(64, 16, stride=1, expand_ratio=6, with_cp=True)
stride=1,
expand_ratio=6,
with_cp=True)
assert block.with_cp assert block.with_cp
x = torch.randn(1, 64, 56, 56) x = torch.randn(1, 64, 56, 56)
x_out = block(x) x_out = block(x)
assert x_out.shape == torch.Size([1, 16, 56, 56]) assert x_out.shape == torch.Size([1, 16, 56, 56])
# Test InvertedResidual with activation=nn.ReLU # Test InvertedResidual with activation=nn.ReLU
block = InvertedResidual(64, 16, block = InvertedResidual(
stride=1, 64, 16, stride=1, expand_ratio=6, activation=nn.ReLU)
expand_ratio=6,
activation=nn.ReLU)
x = torch.randn(1, 64, 56, 56) x = torch.randn(1, 64, 56, 56)
x_out = block(x) x_out = block(x)
assert x_out.shape == torch.Size([1, 16, 56, 56]) assert x_out.shape == torch.Size([1, 16, 56, 56])
@ -193,8 +183,8 @@ def test_mobilenetv2_backbone():
assert feat[6].shape == torch.Size([1, 320, 7, 7]) assert feat[6].shape == torch.Size([1, 320, 7, 7])
# Test MobileNetv2 with layers 1, 3, 5 out forward # Test MobileNetv2 with layers 1, 3, 5 out forward
model = MobileNetv2(widen_factor=1.0, activation=nn.ReLU6, model = MobileNetv2(
out_indices=(0, 2, 4)) widen_factor=1.0, activation=nn.ReLU6, out_indices=(0, 2, 4))
model.init_weights() model.init_weights()
model.train() model.train()
@ -206,8 +196,7 @@ def test_mobilenetv2_backbone():
assert feat[2].shape == torch.Size([1, 96, 14, 14]) assert feat[2].shape == torch.Size([1, 96, 14, 14])
# Test MobileNetv2 with checkpoint forward # Test MobileNetv2 with checkpoint forward
model = MobileNetv2(widen_factor=1.0, activation=nn.ReLU6, model = MobileNetv2(widen_factor=1.0, activation=nn.ReLU6, with_cp=True)
with_cp=True)
for m in model.modules(): for m in model.modules():
if is_block(m): if is_block(m):
assert m.with_cp assert m.with_cp