Huawei_Technology
/
mindspore
Not watched
1
0
Fork 0
Code
Releases 13 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: a146f982bc
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'a146f982bc'
${ noResults }
mindspore/model_zoo/official/cv/yolov4/src/cspdarknet53.py

221 lines
7.8 kB
Raw Blame History 

  1. # Copyright 2020 Huawei Technologies Co., Ltd

  2. #

  3. # Licensed under the Apache License, Version 2.0 (the "License");

  4. # you may not use this file except in compliance with the License.

  5. # You may obtain a copy of the License at

  6. #

  7. # http://www.apache.org/licenses/LICENSE-2.0

  8. #

  9. # Unless required by applicable law or agreed to in writing, software

  10. # distributed under the License is distributed on an "AS IS" BASIS,

  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

  12. # See the License for the specific language governing permissions and

  13. # limitations under the License.

  14. # ============================================================================

  15. """DarkNet model."""

  16. import mindspore.nn as nn

  17. from mindspore.ops import operations as P

  18. 

  19. 

  20. class Mish(nn.Cell):

  21.     """Mish activation method"""

  22.     def __init__(self):

  23.         super(Mish, self).__init__()

  24.         self.mul = P.Mul()

  25.         self.tanh = P.Tanh()

  26.         self.softplus = P.Softplus()

  27. 

  28.     def construct(self, input_x):

  29.         res1 = self.softplus(input_x)

  30.         tanh = self.tanh(res1)

  31.         output = self.mul(input_x, tanh)

  32. 

  33.         return output

  34. 

  35. def conv_block(in_channels,

  36.                out_channels,

  37.                kernel_size,

  38.                stride,

  39.                dilation=1):

  40.     """Get a conv2d batchnorm and relu layer"""

  41.     pad_mode = 'same'

  42.     padding = 0

  43. 

  44.     return nn.SequentialCell(

  45.         [nn.Conv2d(in_channels,

  46.                    out_channels,

  47.                    kernel_size=kernel_size,

  48.                    stride=stride,

  49.                    padding=padding,

  50.                    dilation=dilation,

  51.                    pad_mode=pad_mode),

  52.          nn.BatchNorm2d(out_channels, momentum=0.9, eps=1e-5),

  53.          Mish()

  54.          ]

  55.     )

  56. 

  57. 

  58. class ResidualBlock(nn.Cell):

  59.     """

  60.     DarkNet V1 residual block definition.

  61. 

  62.     Args:

  63.         in_channels: Integer. Input channel.

  64.         out_channels: Integer. Output channel.

  65. 

  66.     Returns:

  67.         Tensor, output tensor.

  68.     Examples:

  69.         ResidualBlock(3, 208)

  70.     """

  71.     def __init__(self,

  72.                  in_channels,

  73.                  out_channels):

  74. 

  75.         super(ResidualBlock, self).__init__()

  76.         out_chls = out_channels

  77.         self.conv1 = conv_block(in_channels, out_chls, kernel_size=1, stride=1)

  78.         self.conv2 = conv_block(out_chls, out_channels, kernel_size=3, stride=1)

  79.         self.add = P.TensorAdd()

  80. 

  81.     def construct(self, x):

  82.         identity = x

  83.         out = self.conv1(x)

  84.         out = self.conv2(out)

  85.         out = self.add(out, identity)

  86. 

  87.         return out

  88. 

  89. class CspDarkNet53(nn.Cell):

  90.     """

  91.     DarkNet V1 network.

  92. 

  93.     Args:

  94.         block: Cell. Block for network.

  95.         layer_nums: List. Numbers of different layers.

  96.         in_channels: Integer. Input channel.

  97.         out_channels: Integer. Output channel.

  98.         num_classes: Integer. Class number. Default:100.

  99. 

  100.     Returns:

  101.         Tuple, tuple of output tensor,(f1,f2,f3,f4,f5).

  102. 

  103.     Examples:

  104.         DarkNet(ResidualBlock)

  105.     """

  106.     def __init__(self,

  107.                  block,

  108.                  detect=False):

  109.         super(CspDarkNet53, self).__init__()

  110. 

  111.         self.outchannel = 1024

  112.         self.detect = detect

  113.         self.concat = P.Concat(axis=1)

  114.         self.add = P.TensorAdd()

  115. 

  116.         self.conv0 = conv_block(3, 32, kernel_size=3, stride=1)

  117.         self.conv1 = conv_block(32, 64, kernel_size=3, stride=2)

  118.         self.conv2 = conv_block(64, 64, kernel_size=1, stride=1)

  119.         self.conv3 = conv_block(64, 32, kernel_size=1, stride=1)

  120.         self.conv4 = conv_block(32, 64, kernel_size=3, stride=1)

  121.         self.conv5 = conv_block(64, 64, kernel_size=1, stride=1)

  122.         self.conv6 = conv_block(64, 64, kernel_size=1, stride=1)

  123.         self.conv7 = conv_block(128, 64, kernel_size=1, stride=1)

  124.         self.conv8 = conv_block(64, 128, kernel_size=3, stride=2)

  125.         self.conv9 = conv_block(128, 64, kernel_size=1, stride=1)

  126.         self.conv10 = conv_block(64, 64, kernel_size=1, stride=1)

  127.         self.conv11 = conv_block(128, 64, kernel_size=1, stride=1)

  128.         self.conv12 = conv_block(128, 128, kernel_size=1, stride=1)

  129.         self.conv13 = conv_block(128, 256, kernel_size=3, stride=2)

  130.         self.conv14 = conv_block(256, 128, kernel_size=1, stride=1)

  131.         self.conv15 = conv_block(128, 128, kernel_size=1, stride=1)

  132.         self.conv16 = conv_block(256, 128, kernel_size=1, stride=1)

  133.         self.conv17 = conv_block(256, 256, kernel_size=1, stride=1)

  134.         self.conv18 = conv_block(256, 512, kernel_size=3, stride=2)

  135.         self.conv19 = conv_block(512, 256, kernel_size=1, stride=1)

  136.         self.conv20 = conv_block(256, 256, kernel_size=1, stride=1)

  137.         self.conv21 = conv_block(512, 256, kernel_size=1, stride=1)

  138.         self.conv22 = conv_block(512, 512, kernel_size=1, stride=1)

  139.         self.conv23 = conv_block(512, 1024, kernel_size=3, stride=2)

  140.         self.conv24 = conv_block(1024, 512, kernel_size=1, stride=1)

  141.         self.conv25 = conv_block(512, 512, kernel_size=1, stride=1)

  142.         self.conv26 = conv_block(1024, 512, kernel_size=1, stride=1)

  143.         self.conv27 = conv_block(1024, 1024, kernel_size=1, stride=1)

  144. 

  145.         self.layer2 = self._make_layer(block, 2, in_channel=64, out_channel=64)

  146.         self.layer3 = self._make_layer(block, 8, in_channel=128, out_channel=128)

  147.         self.layer4 = self._make_layer(block, 8, in_channel=256, out_channel=256)

  148.         self.layer5 = self._make_layer(block, 4, in_channel=512, out_channel=512)

  149. 

  150.     def _make_layer(self, block, layer_num, in_channel, out_channel):

  151.         """

  152.         Make Layer for DarkNet.

  153. 

  154.         :param block: Cell. DarkNet block.

  155.         :param layer_num: Integer. Layer number.

  156.         :param in_channel: Integer. Input channel.

  157.         :param out_channel: Integer. Output channel.

  158.         :return: SequentialCell, the output layer.

  159. 

  160.         Examples:

  161.             _make_layer(ConvBlock, 1, 128, 256)

  162.         """

  163.         layers = []

  164.         darkblk = block(in_channel, out_channel)

  165.         layers.append(darkblk)

  166. 

  167.         for _ in range(1, layer_num):

  168.             darkblk = block(out_channel, out_channel)

  169.             layers.append(darkblk)

  170. 

  171.         return nn.SequentialCell(layers)

  172. 

  173.     def construct(self, x):

  174.         """construct method"""

  175.         c1 = self.conv0(x)

  176.         c2 = self.conv1(c1)  #route

  177.         c3 = self.conv2(c2)

  178.         c4 = self.conv3(c3)

  179.         c5 = self.conv4(c4)

  180.         c6 = self.add(c3, c5)

  181.         c7 = self.conv5(c6)

  182.         c8 = self.conv6(c2)

  183.         c9 = self.concat((c7, c8))

  184.         c10 = self.conv7(c9)

  185.         c11 = self.conv8(c10)   #route

  186.         c12 = self.conv9(c11)

  187.         c13 = self.layer2(c12)

  188.         c14 = self.conv10(c13)

  189.         c15 = self.conv11(c11)

  190.         c16 = self.concat((c14, c15))

  191.         c17 = self.conv12(c16)

  192.         c18 = self.conv13(c17)  #route

  193.         c19 = self.conv14(c18)

  194.         c20 = self.layer3(c19)

  195.         c21 = self.conv15(c20)

  196.         c22 = self.conv16(c18)

  197.         c23 = self.concat((c21, c22))

  198.         c24 = self.conv17(c23)  #output1

  199.         c25 = self.conv18(c24)  #route

  200.         c26 = self.conv19(c25)

  201.         c27 = self.layer4(c26)

  202.         c28 = self.conv20(c27)

  203.         c29 = self.conv21(c25)

  204.         c30 = self.concat((c28, c29))

  205.         c31 = self.conv22(c30)  #output2

  206.         c32 = self.conv23(c31)  #route

  207.         c33 = self.conv24(c32)

  208.         c34 = self.layer5(c33)

  209.         c35 = self.conv25(c34)

  210.         c36 = self.conv26(c32)

  211.         c37 = self.concat((c35, c36))

  212.         c38 = self.conv27(c37)  #output3

  213. 

  214.         if self.detect:

  215.             return c24, c31, c38

  216. 

  217.         return c38

  218. 

  219.     def get_out_channels(self):

  220.         return self.outchannel