109 lines
3.6 KiB
Python
109 lines
3.6 KiB
Python
"""mobile net v2"""
|
|
from mindspore import nn
|
|
from mindspore.ops import operations as P
|
|
|
|
|
|
def make_divisible(input_x, div_by=8):
|
|
return int((input_x + div_by) // div_by)
|
|
|
|
|
|
def _conv_bn(in_channel,
|
|
out_channel,
|
|
ksize,
|
|
stride=1):
|
|
"""Get a conv2d batchnorm and relu layer."""
|
|
return nn.SequentialCell(
|
|
[nn.Conv2dBnAct(in_channel,
|
|
out_channel,
|
|
kernel_size=ksize,
|
|
stride=stride,
|
|
has_bn=True)])
|
|
|
|
|
|
class InvertedResidual(nn.Cell):
|
|
def __init__(self, inp, oup, stride, expend_ratio):
|
|
super(InvertedResidual, self).__init__()
|
|
self.stride = stride
|
|
assert stride in [1, 2]
|
|
|
|
hidden_dim = int(inp * expend_ratio)
|
|
self.use_res_connect = self.stride == 1 and inp == oup
|
|
if expend_ratio == 1:
|
|
self.conv = nn.SequentialCell([
|
|
nn.Conv2dBnAct(hidden_dim,
|
|
hidden_dim,
|
|
3,
|
|
stride,
|
|
group=hidden_dim,
|
|
has_bn=True,
|
|
activation='relu6'),
|
|
nn.Conv2dBnAct(hidden_dim, oup, 1, 1,
|
|
has_bn=True)
|
|
])
|
|
else:
|
|
self.conv = nn.SequentialCell([
|
|
nn.Conv2dBnAct(inp, hidden_dim, 1, 1,
|
|
has_bn=True,
|
|
activation='relu6'),
|
|
nn.Conv2dBnAct(hidden_dim,
|
|
hidden_dim,
|
|
3,
|
|
stride,
|
|
group=hidden_dim,
|
|
has_bn=True,
|
|
activation='relu6'),
|
|
nn.Conv2dBnAct(hidden_dim, oup, 1, 1,
|
|
has_bn=True)
|
|
])
|
|
self.add = P.TensorAdd()
|
|
|
|
def construct(self, input_x):
|
|
out = self.conv(input_x)
|
|
if self.use_res_connect:
|
|
out = self.add(input_x, out)
|
|
return out
|
|
|
|
|
|
class MobileNetV2(nn.Cell):
|
|
def __init__(self, num_class=1000, input_size=224, width_mul=1.):
|
|
super(MobileNetV2, self).__init__()
|
|
_ = input_size
|
|
block = InvertedResidual
|
|
input_channel = 32
|
|
last_channel = 1280
|
|
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, 230, 1, 1],
|
|
]
|
|
if width_mul > 1.0:
|
|
last_channel = make_divisible(last_channel * width_mul)
|
|
self.last_channel = last_channel
|
|
features = [_conv_bn(3, input_channel, 3, 2)]
|
|
|
|
for t, c, n, s in inverted_residual_setting:
|
|
out_channel = make_divisible(c * width_mul) if t > 1 else c
|
|
for i in range(n):
|
|
if i == 0:
|
|
features.append(block(input_channel, out_channel, s, t))
|
|
else:
|
|
features.append(block(input_channel, out_channel, 1, t))
|
|
input_channel = out_channel
|
|
|
|
features.append(_conv_bn(input_channel, self.last_channel, 1))
|
|
|
|
self.features = nn.SequentialCell(features)
|
|
self.mean = P.ReduceMean(keep_dims=False)
|
|
self.classifier = nn.DenseBnAct(self.last_channel, num_class)
|
|
|
|
def construct(self, input_x):
|
|
out = input_x
|
|
out = self.features(out)
|
|
out = self.mean(out, (2, 3))
|
|
out = self.classifier(out)
|
|
return out
|