From ba391c029a4a37159f75a2332585b7a6fd4a30a0 Mon Sep 17 00:00:00 2001 From: louzan Date: Thu, 28 May 2020 11:24:04 +0800 Subject: [PATCH 1/5] add mobilenetv2 --- mmcls/models/backbones/mobilenet_v2.py | 275 +++++++++++++++++++++++++ 1 file changed, 275 insertions(+) create mode 100644 mmcls/models/backbones/mobilenet_v2.py diff --git a/mmcls/models/backbones/mobilenet_v2.py b/mmcls/models/backbones/mobilenet_v2.py new file mode 100644 index 00000000..82e100ae --- /dev/null +++ b/mmcls/models/backbones/mobilenet_v2.py @@ -0,0 +1,275 @@ +import logging + +import torch.nn as nn +import torch.utils.checkpoint as cp + +from ..runner import load_checkpoint +from .base_backbone import BaseBackbone +from .weight_init import constant_init, kaiming_init + + +def conv3x3(in_planes, out_planes, stride=1, dilation=1): + """3x3 convolution with padding""" + return nn.Conv2d( + in_planes, + out_planes, + kernel_size=3, + stride=stride, + padding=dilation, + dilation=dilation, + bias=False) + + +def conv_1x1_bn(inp, oup, act=nn.ReLU6): + return nn.Sequential( + nn.Conv2d(inp, oup, 1, 1, 0, bias=False), + nn.BatchNorm2d(oup), + act(inplace=True) + ) + + +class ConvBNReLU(nn.Sequential): + def __init__(self, + in_planes, + out_planes, + kernel_size=3, + stride=1, + groups=1, + activation=nn.ReLU6): + padding = (kernel_size - 1) // 2 + + try: + self.activation = activation(inplace=True) + except RuntimeWarning('inplace is not allowed to use'): + self.activation = activation() + + super(ConvBNReLU, self).__init__( + nn.Conv2d(in_planes, + out_planes, + kernel_size, + stride, + padding, + groups=groups, + bias=False), + nn.BatchNorm2d(out_planes), + self.activation + ) + + +def _make_divisible(v, divisor, min_value=None): + if min_value is None: + min_value = divisor + new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) + # Make sure that round down does not go down by more than 10%. + if new_v < 0.9 * v: + new_v += divisor + return new_v + + +class InvertedResidual(nn.Module): + def __init__(self, + inplanes, + outplanes, + stride, + expand_ratio, + activation=nn.ReLU6, + with_cp=False): + super(InvertedResidual, self).__init__() + self.stride = stride + assert stride in [1, 2] + self.with_cp = with_cp + self.use_res_connect = self.stride == 1 and inplanes == outplanes + hidden_dim = int(round(inplanes * expand_ratio)) + + layers = [] + if expand_ratio != 1: + # pw + layers.append(ConvBNReLU(inplanes, + hidden_dim, + kernel_size=1, + activation=activation)) + layers.extend([ + # dw + ConvBNReLU(hidden_dim, + hidden_dim, + stride=stride, + groups=hidden_dim, + activation=activation), + # pw-linear + nn.Conv2d(hidden_dim, outplanes, 1, 1, 0, bias=False), + nn.BatchNorm2d(outplanes), + ]) + self.conv = nn.Sequential(*layers) + + def forward(self, x): + def _inner_forward(x): + if self.use_res_connect: + return x + self.conv(x) + else: + return self.conv(x) + + if self.with_cp and x.requires_grad: + out = cp.checkpoint(_inner_forward, x) + else: + out = _inner_forward(x) + + return out + + +def make_inverted_res_layer(block, + inplanes, + planes, + num_blocks, + stride=1, + expand_ratio=6, + activation_type=nn.ReLU6, + with_cp=False): + layers = [] + for i in range(num_blocks): + if i == 0: + layers.append(block(inplanes, planes, stride, + expand_ratio=expand_ratio, + activation=activation_type, + with_cp=with_cp)) + else: + layers.append(block(inplanes, planes, 1, + expand_ratio=expand_ratio, + activation=activation_type, + with_cp=with_cp)) + return nn.Sequential(*layers) + + +class MobileNetv2(BaseBackbone): + """MobileNetv2 backbone. + + Args: + widen_factor (float): Config of widen_factor. + activation (str): Activation type of the network. + out_indices (Sequence[int]): Output from which stages. + frozen_stages (int): Stages to be frozen (all param fixed). -1 means + not freezing any parameters. + bn_eval (bool): Whether to set BN layers as eval mode, namely, freeze + running stats (mean and var). + bn_frozen (bool): Whether to freeze weight and bias of BN layers. + with_cp (bool): Use checkpoint or not. Using checkpoint will save some + memory while slowing down the training speed. + """ + + def __init__(self, + widen_factor=1., + activation=nn.ReLU6, + out_indices=(0, 1, 2, 3, 4, 5, 6), + frozen_stages=-1, + bn_eval=True, + bn_frozen=False, + with_cp=False): + super(MobileNetv2, self).__init__() + block = InvertedResidual + + inverted_residual_setting = { + # lager_index: [expand_ratio, out_channel, n, stide] + 0: [1, 16, 1, 1], + 1: [6, 24, 2, 2], + 2: [6, 32, 3, 2], + 3: [6, 64, 4, 2], + 4: [6, 96, 3, 1], + 5: [6, 160, 3, 2], + 6: [6, 320, 1, 1] + } + self.widen_factor = widen_factor + self.activation_type = activation + try: + self.activation = activation(inplace=True) + except RuntimeWarning('inplace is not allowed to use'): + self.activation = activation() + + self.out_indices = out_indices + self.frozen_stages = frozen_stages + self.bn_eval = bn_eval + self.bn_frozen = bn_frozen + self.with_cp = with_cp + + self.inplanes = 32 + self.inplanes = _make_divisible(self.inplanes * widen_factor, 8) + self.conv1 = conv3x3(3, self.inplanes, stride=2) + + self.inverted_res_layers = [] + for i, later_cfg in enumerate(inverted_residual_setting): + t, c, n, s = later_cfg + planes = _make_divisible(c * widen_factor, 8) + inverted_res_layer = make_inverted_res_layer( + block, + self.inplanes, + planes, + num_blocks=n, + stride=s, + expand_ratio=t, + activation_type=self.activation_type, + with_cp=self.with_cp) + self.inplanes = planes + layer_name = 'layer{}'.format(i + 1) + self.add_module(layer_name, inverted_res_layer) + self.inverted_res_layers.append(layer_name) + + self.out_channel = 1280 + self.out_channel = int(self.out_channel * widen_factor) \ + if widen_factor > 1.0 else self.out_channel + self.conv1_bn = conv_1x1_bn(self.inplanes, self.out_channel) + + self.feat_dim = self.out_channel + + def init_weights(self, pretrained=None): + if isinstance(pretrained, str): + logger = logging.getLogger() + load_checkpoint(self, pretrained, strict=False, logger=logger) + elif pretrained is None: + for m in self.modules(): + if isinstance(m, nn.Conv2d): + kaiming_init(m) + elif isinstance(m, nn.BatchNorm2d): + constant_init(m, 1) + else: + raise TypeError('pretrained must be a str or None') + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.activation(x) + + outs = [] + for i, layer_name in enumerate(self.inverted_res_layers): + inverted_res_layer = getattr(self, layer_name) + x = inverted_res_layer(x) + if i in self.out_indices: + outs.append(x) + + x = self.conv1_bn(x) + outs.append(x) + + if len(outs) == 1: + return outs[0] + else: + return tuple(outs) + + def train(self, mode=True): + super(MobileNetv2, self).train(mode) + if self.bn_eval: + for m in self.modules(): + if isinstance(m, nn.BatchNorm2d): + m.eval() + if self.bn_frozen: + for params in m.parameters(): + params.requires_grad = False + if mode and self.frozen_stages >= 0: + for param in self.conv1.parameters(): + param.requires_grad = False + for param in self.bn1.parameters(): + param.requires_grad = False + self.bn1.eval() + self.bn1.weight.requires_grad = False + self.bn1.bias.requires_grad = False + for i in range(1, self.frozen_stages + 1): + mod = getattr(self, 'layer{}'.format(i)) + mod.eval() + for param in mod.parameters(): + param.requires_grad = False From 2aaff0e4f2725e94659239ba5535267965b405dd Mon Sep 17 00:00:00 2001 From: lixiaojie Date: Wed, 3 Jun 2020 15:51:17 +0800 Subject: [PATCH 2/5] fix & add test --- mmcls/models/backbones/__init__.py | 5 ++ mmcls/models/backbones/base_backbone.py | 27 ++++++++++ mmcls/models/backbones/mobilenet_v2.py | 69 ++++++++++++------------- mmcls/models/backbones/weight_init.py | 66 +++++++++++++++++++++++ tests/test_backbone.py | 25 +++++++++ 5 files changed, 157 insertions(+), 35 deletions(-) create mode 100644 mmcls/models/backbones/base_backbone.py create mode 100644 mmcls/models/backbones/weight_init.py create mode 100644 tests/test_backbone.py diff --git a/mmcls/models/backbones/__init__.py b/mmcls/models/backbones/__init__.py index e69de29b..f66558e0 100644 --- a/mmcls/models/backbones/__init__.py +++ b/mmcls/models/backbones/__init__.py @@ -0,0 +1,5 @@ +from .mobilenet_v2 import MobileNetv2 + +__all__ = [ + 'MobileNetv2', +] \ No newline at end of file diff --git a/mmcls/models/backbones/base_backbone.py b/mmcls/models/backbones/base_backbone.py new file mode 100644 index 00000000..703e0f59 --- /dev/null +++ b/mmcls/models/backbones/base_backbone.py @@ -0,0 +1,27 @@ +import logging +import torch.nn as nn + +from abc import ABCMeta, abstractmethod +from mmcv.runner import load_checkpoint + + +class BaseBackbone(nn.Module, metaclass=ABCMeta): + + def __init__(self): + super(BaseBackbone, self).__init__() + + def init_weights(self, pretrained=None): + if isinstance(pretrained, str): + logger = logging.getLogger() + load_checkpoint(self, pretrained, strict=False, logger=logger) + elif pretrained is None: + pass + else: + raise TypeError('pretrained must be a str or None') + + @abstractmethod + def forward(self, x): + pass + + def train(self, mode=True): + super(BaseBackbone, self).train(mode) diff --git a/mmcls/models/backbones/mobilenet_v2.py b/mmcls/models/backbones/mobilenet_v2.py index 82e100ae..c2c45d0c 100644 --- a/mmcls/models/backbones/mobilenet_v2.py +++ b/mmcls/models/backbones/mobilenet_v2.py @@ -3,7 +3,7 @@ import logging import torch.nn as nn import torch.utils.checkpoint as cp -from ..runner import load_checkpoint +# from ..runner import load_checkpoint from .base_backbone import BaseBackbone from .weight_init import constant_init, kaiming_init @@ -20,11 +20,11 @@ def conv3x3(in_planes, out_planes, stride=1, dilation=1): bias=False) -def conv_1x1_bn(inp, oup, act=nn.ReLU6): +def conv_1x1_bn(inp, oup, activation=nn.ReLU6): return nn.Sequential( nn.Conv2d(inp, oup, 1, 1, 0, bias=False), nn.BatchNorm2d(oup), - act(inplace=True) + activation(inplace=True) ) @@ -38,11 +38,6 @@ class ConvBNReLU(nn.Sequential): activation=nn.ReLU6): padding = (kernel_size - 1) // 2 - try: - self.activation = activation(inplace=True) - except RuntimeWarning('inplace is not allowed to use'): - self.activation = activation() - super(ConvBNReLU, self).__init__( nn.Conv2d(in_planes, out_planes, @@ -52,7 +47,7 @@ class ConvBNReLU(nn.Sequential): groups=groups, bias=False), nn.BatchNorm2d(out_planes), - self.activation + activation(inplace=True) ) @@ -122,20 +117,21 @@ def make_inverted_res_layer(block, num_blocks, stride=1, expand_ratio=6, - activation_type=nn.ReLU6, + activation=nn.ReLU6, with_cp=False): layers = [] for i in range(num_blocks): if i == 0: layers.append(block(inplanes, planes, stride, expand_ratio=expand_ratio, - activation=activation_type, + activation=activation, with_cp=with_cp)) else: layers.append(block(inplanes, planes, 1, expand_ratio=expand_ratio, - activation=activation_type, + activation=activation, with_cp=with_cp)) + inplanes = planes return nn.Sequential(*layers) @@ -165,23 +161,20 @@ class MobileNetv2(BaseBackbone): with_cp=False): super(MobileNetv2, self).__init__() block = InvertedResidual - - inverted_residual_setting = { - # lager_index: [expand_ratio, out_channel, n, stide] - 0: [1, 16, 1, 1], - 1: [6, 24, 2, 2], - 2: [6, 32, 3, 2], - 3: [6, 64, 4, 2], - 4: [6, 96, 3, 1], - 5: [6, 160, 3, 2], - 6: [6, 320, 1, 1] - } + # expand_ratio, out_channel, n, stride + inverted_residual_setting = [ + [1, 16, 1, 1], + [6, 24, 2, 2], + [6, 32, 3, 2], + [6, 64, 4, 2], + [6, 96, 3, 1], + [6, 160, 3, 2], + [6, 320, 1, 1] + ] self.widen_factor = widen_factor - self.activation_type = activation - try: - self.activation = activation(inplace=True) - except RuntimeWarning('inplace is not allowed to use'): - self.activation = activation() + if isinstance(activation, str): + activation = eval(activation) + self.activation = activation(inplace=True) self.out_indices = out_indices self.frozen_stages = frozen_stages @@ -191,11 +184,13 @@ class MobileNetv2(BaseBackbone): self.inplanes = 32 self.inplanes = _make_divisible(self.inplanes * widen_factor, 8) - self.conv1 = conv3x3(3, self.inplanes, stride=2) + self.conv1 = conv3x3(3, self.inplanes, stride=2) + self.bn1 = nn.BatchNorm2d(self.inplanes) self.inverted_res_layers = [] - for i, later_cfg in enumerate(inverted_residual_setting): - t, c, n, s = later_cfg + + for i, layer_cfg in enumerate(inverted_residual_setting): + t, c, n, s = layer_cfg planes = _make_divisible(c * widen_factor, 8) inverted_res_layer = make_inverted_res_layer( block, @@ -204,7 +199,7 @@ class MobileNetv2(BaseBackbone): num_blocks=n, stride=s, expand_ratio=t, - activation_type=self.activation_type, + activation=activation, with_cp=self.with_cp) self.inplanes = planes layer_name = 'layer{}'.format(i + 1) @@ -214,7 +209,9 @@ class MobileNetv2(BaseBackbone): self.out_channel = 1280 self.out_channel = int(self.out_channel * widen_factor) \ if widen_factor > 1.0 else self.out_channel - self.conv1_bn = conv_1x1_bn(self.inplanes, self.out_channel) + + self.conv_last = nn.Conv2d(self.inplanes, self.out_channel, 1, 1, 0, bias=False) + self.bn_last = nn.BatchNorm2d(self.out_channel) self.feat_dim = self.out_channel @@ -233,7 +230,6 @@ class MobileNetv2(BaseBackbone): def forward(self, x): x = self.conv1(x) - x = self.bn1(x) x = self.activation(x) outs = [] @@ -243,7 +239,10 @@ class MobileNetv2(BaseBackbone): if i in self.out_indices: outs.append(x) - x = self.conv1_bn(x) + x = self.conv_last(x) + x = self.bn_last(x) + x = self.activation(x) + outs.append(x) if len(outs) == 1: diff --git a/mmcls/models/backbones/weight_init.py b/mmcls/models/backbones/weight_init.py new file mode 100644 index 00000000..e06e6cca --- /dev/null +++ b/mmcls/models/backbones/weight_init.py @@ -0,0 +1,66 @@ +# Copyright (c) Open-MMLab. All rights reserved. +import numpy as np +import torch.nn as nn + + +def constant_init(module, val, bias=0): + if hasattr(module, 'weight') and module.weight is not None: + nn.init.constant_(module.weight, val) + if hasattr(module, 'bias') and module.bias is not None: + nn.init.constant_(module.bias, bias) + + +def xavier_init(module, gain=1, bias=0, distribution='normal'): + assert distribution in ['uniform', 'normal'] + if distribution == 'uniform': + nn.init.xavier_uniform_(module.weight, gain=gain) + else: + nn.init.xavier_normal_(module.weight, gain=gain) + if hasattr(module, 'bias') and module.bias is not None: + nn.init.constant_(module.bias, bias) + + +def normal_init(module, mean=0, std=1, bias=0): + nn.init.normal_(module.weight, mean, std) + if hasattr(module, 'bias') and module.bias is not None: + nn.init.constant_(module.bias, bias) + + +def uniform_init(module, a=0, b=1, bias=0): + nn.init.uniform_(module.weight, a, b) + if hasattr(module, 'bias') and module.bias is not None: + nn.init.constant_(module.bias, bias) + + +def kaiming_init(module, + a=0, + mode='fan_out', + nonlinearity='relu', + bias=0, + distribution='normal'): + assert distribution in ['uniform', 'normal'] + if distribution == 'uniform': + nn.init.kaiming_uniform_( + module.weight, a=a, mode=mode, nonlinearity=nonlinearity) + else: + nn.init.kaiming_normal_( + module.weight, a=a, mode=mode, nonlinearity=nonlinearity) + if hasattr(module, 'bias') and module.bias is not None: + nn.init.constant_(module.bias, bias) + + +def caffe2_xavier_init(module, bias=0): + # `XavierFill` in Caffe2 corresponds to `kaiming_uniform_` in PyTorch + # Acknowledgment to FAIR's internal code + kaiming_init( + module, + a=1, + mode='fan_in', + nonlinearity='leaky_relu', + distribution='uniform') + + +def bias_init_with_prob(prior_prob): + """ initialize conv/fc bias value according to giving probablity""" + bias_init = float(-np.log((1 - prior_prob) / prior_prob)) + return bias_init diff --git a/tests/test_backbone.py b/tests/test_backbone.py new file mode 100644 index 00000000..db7bce95 --- /dev/null +++ b/tests/test_backbone.py @@ -0,0 +1,25 @@ +import pytest +import torch +import torch.nn as nn +from torch.nn.modules import AvgPool2d, GroupNorm +from torch.nn.modules.batchnorm import _BatchNorm + +from mmcls.models.backbones import MobileNetv2 + + +def test_mobilenetv2_backbone(): + # Test MobileNetv2 with widen_factor 1.0, activation nn.ReLU6 + model = MobileNetv2(widen_factor=1.0, activation=nn.ReLU6) + model.init_weights() + model.train() + + imgs = torch.randn(1, 3, 224, 224) + feat = model(imgs) + assert len(feat) == 8 + 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]) From 83cc9fc2792d4ccb875624c1705e44a42241a229 Mon Sep 17 00:00:00 2001 From: lixiaojie Date: Wed, 3 Jun 2020 15:54:37 +0800 Subject: [PATCH 3/5] fix flake8 --- mmcls/models/backbones/__init__.py | 2 +- mmcls/models/backbones/mobilenet_v2.py | 6 ++++-- tests/test_backbone.py | 3 --- 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/mmcls/models/backbones/__init__.py b/mmcls/models/backbones/__init__.py index f66558e0..1de5cc09 100644 --- a/mmcls/models/backbones/__init__.py +++ b/mmcls/models/backbones/__init__.py @@ -2,4 +2,4 @@ from .mobilenet_v2 import MobileNetv2 __all__ = [ 'MobileNetv2', -] \ No newline at end of file +] diff --git a/mmcls/models/backbones/mobilenet_v2.py b/mmcls/models/backbones/mobilenet_v2.py index c2c45d0c..57731cf8 100644 --- a/mmcls/models/backbones/mobilenet_v2.py +++ b/mmcls/models/backbones/mobilenet_v2.py @@ -3,7 +3,7 @@ import logging import torch.nn as nn import torch.utils.checkpoint as cp -# from ..runner import load_checkpoint +from ..runner import load_checkpoint from .base_backbone import BaseBackbone from .weight_init import constant_init, kaiming_init @@ -210,7 +210,9 @@ class MobileNetv2(BaseBackbone): self.out_channel = int(self.out_channel * widen_factor) \ if widen_factor > 1.0 else self.out_channel - self.conv_last = nn.Conv2d(self.inplanes, self.out_channel, 1, 1, 0, bias=False) + self.conv_last = nn.Conv2d(self.inplanes, + self.out_channel, + 1, 1, 0, bias=False) self.bn_last = nn.BatchNorm2d(self.out_channel) self.feat_dim = self.out_channel diff --git a/tests/test_backbone.py b/tests/test_backbone.py index db7bce95..0e781922 100644 --- a/tests/test_backbone.py +++ b/tests/test_backbone.py @@ -1,8 +1,5 @@ -import pytest import torch import torch.nn as nn -from torch.nn.modules import AvgPool2d, GroupNorm -from torch.nn.modules.batchnorm import _BatchNorm from mmcls.models.backbones import MobileNetv2 From 9712d18f9c5e9660fe7c92675646eddd80d1bd83 Mon Sep 17 00:00:00 2001 From: lixiaojie Date: Wed, 3 Jun 2020 15:57:07 +0800 Subject: [PATCH 4/5] sort --- mmcls/models/backbones/base_backbone.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mmcls/models/backbones/base_backbone.py b/mmcls/models/backbones/base_backbone.py index 703e0f59..1884bd49 100644 --- a/mmcls/models/backbones/base_backbone.py +++ b/mmcls/models/backbones/base_backbone.py @@ -1,7 +1,8 @@ import logging +from abc import ABCMeta, abstractmethod + import torch.nn as nn -from abc import ABCMeta, abstractmethod from mmcv.runner import load_checkpoint From e6987f1c3bd5ea8509751eeed77dbb469422750f Mon Sep 17 00:00:00 2001 From: lixiaojie Date: Wed, 3 Jun 2020 15:58:45 +0800 Subject: [PATCH 5/5] fix import sort --- mmcls/models/backbones/base_backbone.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mmcls/models/backbones/base_backbone.py b/mmcls/models/backbones/base_backbone.py index 1884bd49..65dee16e 100644 --- a/mmcls/models/backbones/base_backbone.py +++ b/mmcls/models/backbones/base_backbone.py @@ -2,7 +2,6 @@ import logging from abc import ABCMeta, abstractmethod import torch.nn as nn - from mmcv.runner import load_checkpoint