|
- # Copyright 2020 Huawei Technologies Co., Ltd
- #
- # Licensed under the Apache License, Version 2.0 (the "License");
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
- # ============================================================================
- """ResNet."""
- import numpy as np
- from scipy.stats import truncnorm
- import mindspore as ms
- import mindspore.nn as nn
- from mindspore.ops import operations as P
- from mindspore.common.tensor import Tensor
-
- format_ = "NHWC"
- # tranpose shape to NCHW, default init is NHWC.
- def _trans_shape(shape, shape_format):
- if shape_format == "NCHW":
- return (shape[0], shape[3], shape[1], shape[2])
- return shape
-
- def _conv_variance_scaling_initializer(in_channel, out_channel, kernel_size):
- fan_in = in_channel * kernel_size * kernel_size
- scale = 1.0
- scale /= max(1., fan_in)
- stddev = (scale ** 0.5) / .87962566103423978
- mu, sigma = 0, stddev
- weight = truncnorm(-2, 2, loc=mu, scale=sigma).rvs(out_channel * in_channel * kernel_size * kernel_size)
- weight = np.reshape(weight, (out_channel, kernel_size, kernel_size, in_channel))
- return Tensor(weight, dtype=ms.float32)
-
- def _weight_variable(shape, factor=0.01):
- init_value = np.random.randn(*shape).astype(np.float32) * factor
- return Tensor(init_value)
-
-
- def _conv3x3(in_channel, out_channel, stride=1):
- weight_shape = (out_channel, 3, 3, in_channel)
- weight_shape = _trans_shape(weight_shape, format_)
- weight = _weight_variable(weight_shape)
- return nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=stride,
- padding=1, pad_mode='pad', weight_init=weight, data_format=format_)
-
- def _conv1x1(in_channel, out_channel, stride=1):
- weight_shape = (out_channel, 1, 1, in_channel)
- weight_shape = _trans_shape(weight_shape, format_)
- weight = _weight_variable(weight_shape)
- return nn.Conv2d(in_channel, out_channel, kernel_size=1, stride=stride,
- padding=0, pad_mode='pad', weight_init=weight, data_format=format_)
-
- def _conv7x7(in_channel, out_channel, stride=1):
- weight_shape = (out_channel, 7, 7, in_channel)
- weight_shape = _trans_shape(weight_shape, format_)
- weight = _weight_variable(weight_shape)
- return nn.Conv2d(in_channel, out_channel, kernel_size=7, stride=stride,
- padding=3, pad_mode='pad', weight_init=weight, data_format=format_)
-
-
- def _bn(channel):
- return nn.BatchNorm2d(channel, eps=1e-4, momentum=0.9, gamma_init=1, beta_init=0,
- moving_mean_init=0, moving_var_init=1, data_format=format_)
-
- def _bn_last(channel):
- return nn.BatchNorm2d(channel, eps=1e-4, momentum=0.9, gamma_init=0, beta_init=0,
- moving_mean_init=0, moving_var_init=1, data_format=format_)
-
- def _fc(in_channel, out_channel):
- weight_shape = (out_channel, in_channel)
- weight = _weight_variable(weight_shape)
- return nn.Dense(in_channel, out_channel, has_bias=True, weight_init=weight, bias_init=0)
-
-
- class ResidualBlock(nn.Cell):
- """
- ResNet V1 residual block definition.
-
- Args:
- in_channel (int): Input channel.
- out_channel (int): Output channel.
- stride (int): Stride size for the first convolutional layer. Default: 1.
-
- Returns:
- Tensor, output tensor.
-
- Examples:
- >>> ResidualBlock(3, 256, stride=2)
- """
- expansion = 4
-
- def __init__(self,
- in_channel,
- out_channel,
- stride=1):
- super(ResidualBlock, self).__init__()
- self.stride = stride
- channel = out_channel // self.expansion
- self.conv1 = _conv1x1(in_channel, channel, stride=1)
- self.bn1 = _bn(channel)
- self.conv2 = _conv3x3(channel, channel, stride=stride)
- self.bn2 = _bn(channel)
-
- self.conv3 = _conv1x1(channel, out_channel, stride=1)
- self.bn3 = _bn_last(out_channel)
- self.relu = nn.ReLU()
-
- self.down_sample = False
-
- if stride != 1 or in_channel != out_channel:
- self.down_sample = True
- self.down_sample_layer = None
-
- if self.down_sample:
- self.down_sample_layer = nn.SequentialCell([_conv1x1(in_channel, out_channel, stride), _bn(out_channel)])
- self.add = P.Add()
-
- def construct(self, x):
- identity = x
- if self.down_sample:
- identity = self.down_sample_layer(identity)
-
- out = self.conv1(x)
- out = self.bn1(out)
- out = self.relu(out)
- out = self.conv2(out)
- out = self.bn2(out)
- out = self.relu(out)
- out = self.conv3(out)
- out = self.bn3(out)
-
- out = self.add(identity, out)
- out = self.relu(out)
-
- return out
-
-
- class ResNet(nn.Cell):
- """
- ResNet architecture.
-
- Args:
- block (Cell): Block for network.
- layer_nums (list): Numbers of block in different layers.
- in_channels (list): Input channel in each layer.
- out_channels (list): Output channel in each layer.
- strides (list): Stride size in each layer.
- num_classes (int): The number of classes that the training images are belonging to.
- Returns:
- Tensor, output tensor.
-
- Examples:
- >>> ResNet(ResidualBlock,
- >>> [3, 4, 6, 3],
- >>> [64, 256, 512, 1024],
- >>> [256, 512, 1024, 2048],
- >>> [1, 2, 2, 2],
- >>> 10)
- """
-
- def __init__(self,
- block,
- layer_nums,
- in_channels,
- out_channels,
- strides,
- num_classes):
- super(ResNet, self).__init__()
-
- if not len(layer_nums) == len(in_channels) == len(out_channels) == 4:
- raise ValueError("the length of layer_num, in_channels, out_channels list must be 4!")
- input_data_channel = 4
- if format_ == "NCHW":
- input_data_channel = 3
- self.conv1 = _conv7x7(input_data_channel, 64, stride=2)
- self.bn1 = _bn(64)
- self.relu = P.ReLU()
- self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, pad_mode="same", data_format=format_)
- self.layer1 = self._make_layer(block,
- layer_nums[0],
- in_channel=in_channels[0],
- out_channel=out_channels[0],
- stride=strides[0])
- self.layer2 = self._make_layer(block,
- layer_nums[1],
- in_channel=in_channels[1],
- out_channel=out_channels[1],
- stride=strides[1])
- self.layer3 = self._make_layer(block,
- layer_nums[2],
- in_channel=in_channels[2],
- out_channel=out_channels[2],
- stride=strides[2])
- self.layer4 = self._make_layer(block,
- layer_nums[3],
- in_channel=in_channels[3],
- out_channel=out_channels[3],
- stride=strides[3])
-
- self.avg_pool = P.AvgPool(7, 1, data_format=format_)
- self.flatten = nn.Flatten()
- self.end_point = _fc(out_channels[3], num_classes)
-
- def _make_layer(self, block, layer_num, in_channel, out_channel, stride):
- """
- Make stage network of ResNet.
-
- Args:
- block (Cell): Resnet block.
- layer_num (int): Layer number.
- in_channel (int): Input channel.
- out_channel (int): Output channel.
- stride (int): Stride size for the first convolutional layer.
- Returns:
- SequentialCell, the output layer.
-
- Examples:
- >>> _make_layer(ResidualBlock, 3, 128, 256, 2)
- """
- layers = []
-
- resnet_block = block(in_channel, out_channel, stride=stride)
- layers.append(resnet_block)
- for _ in range(1, layer_num):
- resnet_block = block(out_channel, out_channel, stride=1)
- layers.append(resnet_block)
- return nn.SequentialCell(layers)
-
- def construct(self, x):
- x = self.conv1(x)
- x = self.bn1(x)
- x = self.relu(x)
- c1 = self.maxpool(x)
-
- c2 = self.layer1(c1)
- c3 = self.layer2(c2)
- c4 = self.layer3(c3)
- c5 = self.layer4(c4)
-
- out = self.avg_pool(c5)
- out = self.flatten(out)
- out = self.end_point(out)
-
- return out
-
-
- def resnet50(class_num=1001, dtype="fp16"):
- """
- Get ResNet50 neural network.
-
- Args:
- class_num (int): Class number.
-
- Returns:
- Cell, cell instance of ResNet50 neural network.
-
- Examples:
- >>> net = resnet50(1001)
- """
- global format_
- if dtype == "fp32":
- format_ = "NCHW"
- return ResNet(ResidualBlock,
- [3, 4, 6, 3],
- [64, 256, 512, 1024],
- [256, 512, 1024, 2048],
- [1, 2, 2, 2],
- class_num)
|