|
- import mindspore as ms
- import mindspore.nn as nn
- import mindspore.numpy as np
- import config as cfg
-
-
- class MaxPool(nn.Cell):
- def __init__(self, kernel_size=1, stride=1, padding=0):
- super(MaxPool, self).__init__()
- self.kernel_size = kernel_size
- self.stride = stride
- self.pad = nn.Pad(paddings=((0, 0), (0, 0), (padding, padding), (padding, padding)))
- self.pool = nn.MaxPool2d(self.kernel_size, self.stride)
-
- def construct(self, x):
- x = self.pad(x)
- x = self.pool(x)
- return x
-
-
- class BottleNeck(nn.Cell):
- def __init__(self, in_channel, out_channel, down_sample=False):
- super(BottleNeck, self).__init__()
- self.in_channel = in_channel
- self.out_channel = out_channel
- self.down_sample = down_sample
- self.relu = nn.ReLU()
- if down_sample:
- self.conv1 = nn.Conv2d(in_channel, int(out_channel / 4), 1, 2, 'valid')
- else:
- self.conv1 = nn.Conv2d(in_channel, int(out_channel / 4), 1, 1, 'valid')
- self.conv2 = nn.Conv2d(int(out_channel / 4), int(out_channel / 4), 3, 1, 'same')
- self.conv3 = nn.Conv2d(int(out_channel / 4), out_channel, 1, 1, 'same')
- if down_sample:
- self.residual = nn.Conv2d(in_channel, out_channel, 1, 2, 'valid')
- elif self.in_channel != self.out_channel:
- self.residual = nn.Conv2d(in_channel, out_channel, 1, 1, 'valid')
- self.bn1 = nn.BatchNorm2d(int(out_channel / 4))
- self.bn2 = nn.BatchNorm2d(int(out_channel / 4))
- self.bn3 = nn.BatchNorm2d(out_channel)
- if self.down_sample or (self.in_channel != self.out_channel):
- self.bn4 = nn.BatchNorm2d(out_channel)
-
- def construct(self, x):
- x1 = x
- x = self.conv1(x)
- x = self.bn1(x)
- x = self.relu(x)
- x = self.conv2(x)
- x = self.bn2(x)
- x = self.relu(x)
- x = self.conv3(x)
- x = self.bn3(x)
- if (not self.down_sample) and (self.in_channel == self.out_channel):
- x += x1
- else:
- x1 = self.residual(x1)
- x1 = self.bn4(x1)
- x += x1
- x = self.relu(x)
- return x
-
-
- class Resnet50(nn.Cell):
- def __init__(self):
- super(Resnet50, self).__init__()
- self.relu = nn.ReLU()
- self.conv1 = nn.Conv2d(3, 64, 7, 2, 'pad', 3)
- self.pool1 = MaxPool(3, 2, 1)
- self.bn1 = nn.BatchNorm2d(64)
- self.stage1 = nn.CellList([BottleNeck(64, 128, False)])
- self.stage1.extend([BottleNeck(128, 128, False) for i in range(2)])
- self.stage2 = nn.CellList([BottleNeck(128, 256, True)])
- self.stage2.extend([BottleNeck(256, 256, False) for i in range(3)])
- self.stage3 = nn.CellList([BottleNeck(256, 512, True)])
- self.stage3.extend([BottleNeck(512, 512, False) for i in range(5)])
- self.stage4 = nn.CellList([BottleNeck(512, 1024, True)])
- self.stage4.extend([BottleNeck(1024, 1024, False) for i in range(2)])
- self.pool2 = nn.AvgPool2d((7, 10))
- self.flatten = nn.Flatten()
- self.fc = nn.Dense(1024, cfg.output_num)
-
- def construct(self, x):
- x = self.conv1(x)
- x = self.bn1(x)
- x = self.relu(x)
- x = self.pool1(x)
- for block in self.stage1:
- x = block(x)
- for block in self.stage2:
- x = block(x)
- for block in self.stage3:
- x = block(x)
- for block in self.stage4:
- x = block(x)
- x = self.pool2(x)
- x = self.flatten(x)
- x = self.fc(x)
- return x
-
-
- if __name__ == '__main__':
- net = Resnet50()
- d = np.randn([10, 3, 224, 300], dtype=ms.float32)
- out = net(d)
- print(out.shape)
|