mirror of https://github.com/Jittor/Jittor
102 lines
4.4 KiB
Python
102 lines
4.4 KiB
Python
# ***************************************************************
|
|
# Copyright (c) 2021 Jittor. All Rights Reserved.
|
|
# Maintainers:
|
|
# Wenyang Zhou <576825820@qq.com>
|
|
# Dun Liang <randonlang@gmail.com>.
|
|
#
|
|
# This file is subject to the terms and conditions defined in
|
|
# file 'LICENSE.txt', which is part of this source code package.
|
|
# ***************************************************************
|
|
# This model is generated by pytorch converter.
|
|
|
|
import jittor as jt
|
|
from jittor import init
|
|
from jittor import nn
|
|
__all__ = ['MobileNetV2', 'mobilenet_v2']
|
|
|
|
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))
|
|
if (new_v < (0.9 * v)):
|
|
new_v += divisor
|
|
return new_v
|
|
|
|
class ConvBNReLU(nn.Sequential):
|
|
|
|
def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):
|
|
padding = ((kernel_size - 1) // 2)
|
|
super(ConvBNReLU, self).__init__(nn.Conv(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False), nn.BatchNorm(out_planes), nn.ReLU6())
|
|
|
|
class InvertedResidual(nn.Module):
|
|
|
|
def __init__(self, inp, oup, stride, expand_ratio):
|
|
super(InvertedResidual, self).__init__()
|
|
self.stride = stride
|
|
assert (stride in [1, 2])
|
|
hidden_dim = int(round((inp * expand_ratio)))
|
|
self.use_res_connect = ((self.stride == 1) and (inp == oup))
|
|
layers = []
|
|
if (expand_ratio != 1):
|
|
layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
|
|
layers.extend([ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim), nn.Conv(hidden_dim, oup, 1, 1, 0, bias=False), nn.BatchNorm(oup)])
|
|
self.conv = nn.Sequential(*layers)
|
|
|
|
def execute(self, x):
|
|
if self.use_res_connect:
|
|
return (x + self.conv(x))
|
|
else:
|
|
return self.conv(x)
|
|
|
|
class MobileNetV2(nn.Module):
|
|
""" MobileNetV2 model architecture.
|
|
|
|
Args:
|
|
|
|
* num_classes: Number of classes. Default: 1000.
|
|
* width_mult: Width multiplier - adjusts number of channels in each layer by this amount. Default: 1.0.
|
|
* init_weights: Defualt: True.
|
|
* inverted_residual_setting: Network structure
|
|
* round_nearest: Round the number of channels in each layer to be a multiple of this number. Set to 1 to turn off rounding. Default: 8.
|
|
* block: Module specifying inverted residual building block for mobilenet. If None, use InvertedResidual instead. Default: None.
|
|
"""
|
|
|
|
def __init__(self, num_classes=1000, width_mult=1.0, inverted_residual_setting=None, round_nearest=8, block=None):
|
|
super(MobileNetV2, self).__init__()
|
|
if (block is None):
|
|
block = InvertedResidual
|
|
input_channel = 32
|
|
last_channel = 1280
|
|
if (inverted_residual_setting is None):
|
|
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]]
|
|
if ((len(inverted_residual_setting) == 0) or (len(inverted_residual_setting[0]) != 4)):
|
|
raise ValueError('inverted_residual_setting should be non-empty or a 4-element list, got {}'.format(inverted_residual_setting))
|
|
input_channel = _make_divisible((input_channel * width_mult), round_nearest)
|
|
self.last_channel = _make_divisible((last_channel * max(1.0, width_mult)), round_nearest)
|
|
features = [ConvBNReLU(3, input_channel, stride=2)]
|
|
for (t, c, n, s) in inverted_residual_setting:
|
|
output_channel = _make_divisible((c * width_mult), round_nearest)
|
|
for i in range(n):
|
|
stride = (s if (i == 0) else 1)
|
|
features.append(block(input_channel, output_channel, stride, expand_ratio=t))
|
|
input_channel = output_channel
|
|
features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1))
|
|
self.features = nn.Sequential(*features)
|
|
self.classifier = nn.Sequential(nn.Dropout(0.2), nn.Linear(self.last_channel, num_classes))
|
|
|
|
def _forward_impl(self, x):
|
|
x = self.features(x)
|
|
x = nn.AdaptiveAvgPool2d(1)(x)
|
|
x = jt.reshape(x, (x.shape[0], -1))
|
|
x = self.classifier(x)
|
|
return x
|
|
|
|
def execute(self, x):
|
|
return self._forward_impl(x)
|
|
|
|
def mobilenet_v2(pretrained=False):
|
|
model = MobileNetV2()
|
|
if pretrained: model.load("jittorhub://mobilenet_v2.pkl")
|
|
return model
|
|
|