forked from huawei/mindspore2022
318 lines
11 KiB
Python
Executable File
318 lines
11 KiB
Python
Executable File
"""EffNet model define"""
|
|
import numpy as np
|
|
import mindspore.nn as nn
|
|
from mindspore.ops import operations as P
|
|
from mindspore.common.initializer import TruncatedNormal
|
|
from mindspore import Tensor
|
|
|
|
|
|
__all__ = ['effnet']
|
|
|
|
|
|
def weight_variable():
|
|
"""weight initial"""
|
|
return TruncatedNormal(0.02)
|
|
|
|
|
|
def _make_divisible(v, divisor, min_value=None):
|
|
"""
|
|
This function is taken from the original tf repo.
|
|
It ensures that all layers have a channel number that is divisible by 8
|
|
It can be seen here:
|
|
https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
|
|
:param v:
|
|
:param divisor:
|
|
iparam min_value:
|
|
:return:
|
|
"""
|
|
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 Swish(nn.Cell):
|
|
def __init__(self):
|
|
super().__init__(Swish)
|
|
self.sigmoid = nn.Sigmoid()
|
|
|
|
def construct(self, x):
|
|
s = self.sigmoid(x)
|
|
m = x*s
|
|
return m
|
|
|
|
class AdaptiveAvgPool(nn.Cell):
|
|
def __init__(self, output_size=None):
|
|
super().__init__(AdaptiveAvgPool)
|
|
self.mean = P.ReduceMean(keep_dims=True)
|
|
self.output_size = output_size
|
|
|
|
def construct(self, x):
|
|
return self.mean(x, (2, 3))
|
|
|
|
class SELayer(nn.Cell):
|
|
"""
|
|
SELayer
|
|
"""
|
|
def __init__(self, channel, reduction=4):
|
|
super().__init__(SELayer)
|
|
reduced_chs = _make_divisible(channel/reduction, 1)
|
|
self.avg_pool = AdaptiveAvgPool(output_size=(1, 1))
|
|
weight = weight_variable()
|
|
self.conv_reduce = nn.Conv2d(
|
|
in_channels=channel, out_channels=reduced_chs, kernel_size=1, has_bias=True, weight_init=weight)
|
|
self.act1 = Swish()
|
|
self.conv_expand = nn.Conv2d(
|
|
in_channels=reduced_chs, out_channels=channel, kernel_size=1, has_bias=True)
|
|
self.act2 = nn.Sigmoid()
|
|
|
|
def construct(self, x):
|
|
o = self.avg_pool(x)
|
|
o = self.conv_reduce(o)
|
|
o = self.act1(o)
|
|
o = self.conv_expand(o)
|
|
o = self.act2(o)
|
|
return x * o
|
|
|
|
class DepthwiseSeparableConv(nn.Cell):
|
|
"""
|
|
DepthwiseSeparableConv
|
|
"""
|
|
def __init__(self, in_chs, out_chs, dw_kernel_size=3,
|
|
stride=1, noskip=False, se_ratio=0.0, drop_connect_rate=0.0):
|
|
super().__init__(DepthwiseSeparableConv)
|
|
assert stride in [1, 2]
|
|
self.has_residual = (stride == 1 and in_chs == out_chs) and not noskip
|
|
self.drop_connect_rate = drop_connect_rate
|
|
|
|
self.conv_dw = nn.Conv2d(in_channels=in_chs, out_channels=in_chs, kernel_size=dw_kernel_size,
|
|
stride=stride, pad_mode="pad", padding=1, has_bias=False, group=in_chs)
|
|
self.bn1 = nn.BatchNorm2d(in_chs, eps=0.001)
|
|
self.act1 = Swish()
|
|
|
|
if se_ratio is not None and se_ratio > 0.:
|
|
self.se = SELayer(in_chs, reduction=se_ratio)
|
|
else:
|
|
print("ERRRRRORRRR -- not prepared for this one\n")
|
|
|
|
self.conv_pw = nn.Conv2d(
|
|
in_channels=in_chs, out_channels=out_chs, kernel_size=1, stride=stride, has_bias=False)
|
|
self.bn2 = nn.BatchNorm2d(out_chs, eps=0.001)
|
|
|
|
def construct(self, x):
|
|
"""
|
|
construct
|
|
"""
|
|
residual = x
|
|
x = self.conv_dw(x)
|
|
x = self.bn1(x)
|
|
x = self.act1(x)
|
|
x = self.se(x)
|
|
x = self.conv_pw(x)
|
|
x = self.bn2(x)
|
|
if self.has_residual:
|
|
x += residual
|
|
return x
|
|
|
|
def conv_3x3_bn(inp, oup, stride):
|
|
weight = weight_variable()
|
|
return nn.SequentialCell([
|
|
nn.Conv2d(in_channels=inp, out_channels=oup, kernel_size=3, stride=stride,
|
|
padding=1, weight_init=weight, has_bias=False, pad_mode='pad'),
|
|
nn.BatchNorm2d(oup, eps=0.001), # , momentum=0.1),
|
|
nn.HSwish()])
|
|
|
|
def conv_1x1_bn(inp, oup):
|
|
weight = weight_variable()
|
|
return nn.SequentialCell([
|
|
nn.Conv2d(in_channels=inp, out_channels=oup, kernel_size=1,
|
|
stride=1, padding=0, weight_init=weight, has_bias=False),
|
|
nn.BatchNorm2d(oup, eps=0.001),
|
|
nn.HSwish()])
|
|
|
|
class InvertedResidual(nn.Cell):
|
|
"""
|
|
InvertedResidual
|
|
"""
|
|
def __init__(self, in_chs, out_chs, kernel_size, stride, padding, expansion, se_ratio):
|
|
super().__init__(InvertedResidual)
|
|
assert stride in [1, 2]
|
|
mid_chs: int = _make_divisible(in_chs * expansion, 1)
|
|
self.has_residual = (in_chs == out_chs and stride == 1)
|
|
self.drop_connect_rate = 0
|
|
self.conv_pw = nn.Conv2d(
|
|
in_channels=in_chs, out_channels=mid_chs, kernel_size=1, stride=1, has_bias=False)
|
|
self.bn1 = nn.BatchNorm2d(mid_chs, eps=0.001) # ,momentum=0.1)
|
|
self.act1 = Swish()
|
|
if stride > 1:
|
|
self.conv_dw = nn.Conv2d(in_channels=mid_chs, out_channels=mid_chs, kernel_size=kernel_size,
|
|
stride=stride, padding=padding, has_bias=False, group=mid_chs, pad_mode='same')
|
|
else:
|
|
self.conv_dw = nn.Conv2d(in_channels=mid_chs, out_channels=mid_chs, kernel_size=kernel_size,
|
|
stride=stride, padding=padding, has_bias=False, group=mid_chs, pad_mode='pad')
|
|
self.bn2 = nn.BatchNorm2d(mid_chs, eps=0.001) # ,momentum=0.1)
|
|
self.act2 = Swish()
|
|
|
|
# Squeeze-and-excitation
|
|
if se_ratio is not None and se_ratio > 0.:
|
|
self.se = SELayer(mid_chs, reduction=se_ratio)
|
|
else:
|
|
print("ERRRRRORRRR -- not prepared for this one\n")
|
|
|
|
# Point-wise linear projection
|
|
self.conv_pwl = nn.Conv2d(
|
|
in_channels=mid_chs, out_channels=out_chs, kernel_size=1, stride=1, has_bias=False)
|
|
self.bn3 = nn.BatchNorm2d(out_chs, eps=0.001) # ,momentum=0.1)
|
|
|
|
def construct(self, x):
|
|
"""
|
|
construct
|
|
"""
|
|
residual = x
|
|
|
|
# Point-wise expansion
|
|
x = self.conv_pw(x)
|
|
x = self.bn1(x)
|
|
x = self.act1(x)
|
|
|
|
# Depth-wise convolution
|
|
x = self.conv_dw(x)
|
|
x = self.bn2(x)
|
|
x = self.act2(x)
|
|
|
|
# Squeeze-and-excitation
|
|
x = self.se(x)
|
|
|
|
# Point-wise linear projection
|
|
x = self.conv_pwl(x)
|
|
x = self.bn3(x)
|
|
|
|
if self.has_residual:
|
|
x += residual
|
|
return x
|
|
|
|
class EfficientNet(nn.Cell):
|
|
"""
|
|
EfficientNet
|
|
"""
|
|
def __init__(self, cfgs, num_classes=1000):
|
|
super().__init__(EfficientNet)
|
|
# setting of inverted residual blocks
|
|
self.cfgs = cfgs
|
|
stem_size = 32
|
|
self.num_classes_ = num_classes
|
|
self.num_features_ = 1280
|
|
|
|
self.conv_stem = nn.Conv2d(
|
|
in_channels=3, out_channels=stem_size, kernel_size=3, stride=2, has_bias=False)
|
|
|
|
self.bn1 = nn.BatchNorm2d(stem_size, eps=0.001) # momentum=0.1)
|
|
self.act1 = Swish()
|
|
in_chs = stem_size
|
|
|
|
layers = [nn.SequentialCell([DepthwiseSeparableConv(in_chs, 16, 3, 1, se_ratio=4)]),
|
|
|
|
nn.SequentialCell([InvertedResidual(16, 24, 3, 2, 0, 6, se_ratio=24),
|
|
InvertedResidual(24, 24, 3, 1, 1, 6, se_ratio=24)]),
|
|
|
|
nn.SequentialCell([InvertedResidual(24, 40, 5, 2, 0, 6, se_ratio=24),
|
|
InvertedResidual(40, 40, 5, 1, 2, 6, se_ratio=24)]),
|
|
|
|
nn.SequentialCell([InvertedResidual(40, 80, 3, 2, 0, 6, se_ratio=24),
|
|
InvertedResidual(
|
|
80, 80, 3, 1, 1, 6, se_ratio=24),
|
|
InvertedResidual(80, 80, 3, 1, 1, 6, se_ratio=24)]),
|
|
|
|
nn.SequentialCell([InvertedResidual(80, 112, 5, 1, 2, 6, se_ratio=24),
|
|
InvertedResidual(
|
|
112, 112, 5, 1, 2, 6, se_ratio=24),
|
|
InvertedResidual(112, 112, 5, 1, 2, 6, se_ratio=24)]),
|
|
|
|
nn.SequentialCell([InvertedResidual(112, 192, 5, 2, 0, 6, se_ratio=24),
|
|
InvertedResidual(
|
|
192, 192, 5, 1, 2, 6, se_ratio=24),
|
|
InvertedResidual(
|
|
192, 192, 5, 1, 2, 6, se_ratio=24),
|
|
InvertedResidual(192, 192, 5, 1, 2, 6, se_ratio=24)]),
|
|
|
|
nn.SequentialCell(
|
|
[InvertedResidual(192, 320, 3, 1, 1, 6, se_ratio=24)])
|
|
]
|
|
self.blocks = nn.SequentialCell(layers)
|
|
|
|
self.conv_head = nn.Conv2d(
|
|
in_channels=320, out_channels=self.num_features_, kernel_size=1)
|
|
self.bn2 = nn.BatchNorm2d(self.num_features_, eps=0.001)
|
|
self.act2 = Swish()
|
|
self.global_pool = AdaptiveAvgPool(output_size=(1, 1))
|
|
self.classifier = nn.Dense(self.num_features_, num_classes)
|
|
|
|
self._initialize_weights()
|
|
|
|
def construct(self, x):
|
|
"""
|
|
construct
|
|
"""
|
|
x = self.conv_stem(x)
|
|
x = self.bn1(x)
|
|
x = self.act1(x)
|
|
x = self.blocks(x)
|
|
x = self.conv_head(x)
|
|
x = self.bn2(x)
|
|
x = self.act2(x)
|
|
x = self.global_pool(x)
|
|
x = P.Reshape()(x, (-1, self.num_features_))
|
|
x = self.classifier(x)
|
|
return x
|
|
|
|
def _initialize_weights(self):
|
|
"""
|
|
_initialize_weights
|
|
"""
|
|
def init_linear_weight(m):
|
|
m.weight.set_data(Tensor(np.random.normal(
|
|
0, 0.01, m.weight.data.shape).astype("float32")))
|
|
if m.bias is not None:
|
|
m.bias.set_data(
|
|
Tensor(np.zeros(m.bias.data.shape, dtype="float32")))
|
|
|
|
for m in self.cells():
|
|
if isinstance(m, nn.Conv2d):
|
|
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
|
|
m.weight.set_data(Tensor(np.random.normal(0, np.sqrt(2. / n),
|
|
m.weight.data.shape).astype("float32")))
|
|
if m.bias is not None:
|
|
m.bias.data.zero_()
|
|
m.weight.requires_grad = True
|
|
elif isinstance(m, nn.BatchNorm2d):
|
|
m.gamma.set_data(
|
|
Tensor(np.ones(m.gamma.data.shape, dtype="float32")))
|
|
m.beta.set_data(
|
|
Tensor(np.zeros(m.beta.data.shape, dtype="float32")))
|
|
elif isinstance(m, nn.Dense):
|
|
init_linear_weight(m)
|
|
|
|
def effnet(**kwargs):
|
|
"""
|
|
Constructs a EfficientNet model
|
|
"""
|
|
cfgs = [
|
|
# k, t, c, SE, HS, s
|
|
[3, 1, 16, 1, 0, 2],
|
|
[3, 4.5, 24, 0, 0, 2],
|
|
[3, 3.67, 24, 0, 0, 1],
|
|
[5, 4, 40, 1, 1, 2],
|
|
[5, 6, 40, 1, 1, 1],
|
|
[5, 6, 40, 1, 1, 1],
|
|
[5, 3, 48, 1, 1, 1],
|
|
[5, 3, 48, 1, 1, 1],
|
|
[5, 6, 96, 1, 1, 2],
|
|
[5, 6, 96, 1, 1, 1],
|
|
[5, 6, 96, 1, 1, 1],
|
|
]
|
|
|
|
return EfficientNet(cfgs, **kwargs)
|