|
- # 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.
- # ============================================================================
-
- """Network."""
- import numpy as np
-
- import mindspore
- from mindspore import nn
- from mindspore.ops import operations as P
- from mindspore import Tensor
-
- # conv_weight_init = 'HeUniform'
-
- # ResNet
- def _weight_variable(shape, factor=0.01):
- """_weight_variable"""
- init_value = np.random.randn(*shape).astype(np.float32) * factor
- return Tensor(init_value)
-
- def _conv3x3(in_channel, out_channel, stride=1):
- """conv3x3"""
- weight_shape = (out_channel, in_channel, 3, 3)
- 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)
-
- def _conv1x1(in_channel, out_channel, stride=1):
- """conv1x1"""
- weight_shape = (out_channel, in_channel, 1, 1)
- 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)
-
- def _conv7x7(in_channel, out_channel, stride=1):
- """conv7x7"""
- weight_shape = (out_channel, in_channel, 7, 7)
- 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)
-
- def _bn(channel):
- """bn"""
- return nn.BatchNorm2d(channel)
-
-
- def _bn_last(channel):
- """bn_last"""
- return nn.BatchNorm2d(channel)
-
-
- def _fc(in_channel, out_channel):
- """fc layer"""
- 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):
- """ResidualBlock"""
- expansion = 4
-
- def __init__(self,
- in_channel,
- out_channel,
- stride=1):
- super().__init__()
-
- 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):
- """construct"""
- identity = x
-
- 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)
-
- if self.down_sample:
- identity = self.down_sample_layer(identity)
-
- out = self.add(out, identity)
- out = self.relu(out)
-
- return out
-
- class ResNet(nn.Cell):
- """
- ResNet architecture, returns last 3 layers outputs
-
- Args:
-
- block (Object): ResidualBlock.
- layer_nums (list): A list indicating how many layers each ResidualBlock has.
- in_channels (list): A list indicating the input channels of each ResidualBlock.
- out_channels (list): A list indicating the output channels of each ResidualBlock.
- strides (list): A list indicating the strides of each ResidualBlock.
- classnum (int): num of classes.
-
- Examples:
- >>> resnet50 = ResNet(ResidualBlock,
- [3, 4, 6, 3],
- [64, 256, 512, 1024],
- [256, 512, 1024, 2048],
- [1, 2, 2, 2],
- class_num)
- """
- def __init__(self,
- block,
- layer_nums,
- in_channels,
- out_channels,
- strides,
- num_classes):
- super().__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!")
-
- self.conv1 = _conv7x7(3, 64, stride=2)
- self.bn1 = _bn(64)
- self.relu = P.ReLU()
-
- self.zeros1 = P.Zeros()
- self.zeros2 = P.Zeros()
- self.concat1 = P.Concat(axis=2)
- self.concat2 = P.Concat(axis=3)
- self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, pad_mode="valid")
-
-
- 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.mean = P.ReduceMean(keep_dims=True)
- 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_layer"""
- 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):
- """construct"""
- x = self.conv1(x)
- x = self.bn1(x)
- x = self.relu(x)
- zeros1 = self.zeros1((x.shape[0], x.shape[1], 1, x.shape[3]), mindspore.float32)
- x = self.concat1((zeros1, x))
- zeros2 = self.zeros2((x.shape[0], x.shape[1], x.shape[2], 1), mindspore.float32)
- x = self.concat2((zeros2, x))
-
- c1 = self.maxpool(x)
-
- c2 = self.layer1(c1)
- c3 = self.layer2(c2)
- c4 = self.layer3(c3)
- c5 = self.layer4(c4)
-
- out = self.mean(c5, (2, 3))
- out = self.flatten(out)
- out = self.end_point(out)
-
- return c3, c4, c5
-
- def resnet50(class_num=10):
- """
- resnet50 model, returns last 3 layers outputs
-
- Args:
- classnum (int): num of classes
-
- Examples:
- >>> backbone = resnet50(1001)
- """
- return ResNet(ResidualBlock,
- [3, 4, 6, 3],
- [64, 256, 512, 1024],
- [256, 512, 1024, 2048],
- [1, 2, 2, 2],
- class_num)
|