|
- # Copyright 2022 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.
- # =======================================================================================
- """Darknet related"""
- import mindspore.nn as nn
-
- from src.network_blocks import BaseConv, CSPLayer, DWConv, Focus, ResLayer, SPPBottleneck
-
-
- class Darknet(nn.Cell):
- """ Darknet for yolox-darknet53 """
- # number of block from dark2 to dark5.
- depth2block = {21: [1, 2, 2, 1], 53: [2, 8, 8, 4]}
-
- def __init__(
- self,
- depth,
- in_channels=3,
- stem_out_channels=32,
- out_features=("dark3", "dark4", "dark5"),
- ):
- """
- Args:
- depth (int): depth of darknet used in model, usually use [21, 53] for this param.
- in_channels (int): number of input channels, for example, use 3 for RGB image.
- stem_out_channels (int): number of output channels of darknet stem.
- It decides channels of darknet layer2 to layer5.
- out_features (Tuple[str]): desired output layer name.
- """
- super(Darknet, self).__init__()
- assert out_features, "please provide output features of Darknet"
- self.out_features = out_features
- self.stem = nn.SequentialCell(
- BaseConv(in_channels=in_channels, out_channels=stem_out_channels, ksize=3, stride=1, act="lrelu"),
- *self.make_group_layer(stem_out_channels, num_blocks=1, stride=2),
- )
- in_channels = stem_out_channels * 2
-
- num_blocks = Darknet.depth2block[depth]
- # create darknet with `stem_out_channels` and `num_blocks` layers.
- # to make model structure more clear, we don't use `for` statement in python.
- self.dark2 = nn.SequentialCell(
- *self.make_group_layer(in_channels=in_channels, num_blocks=num_blocks[0], stride=2)
- )
- in_channels *= 2 # 128
- self.dark3 = nn.SequentialCell(
- *self.make_group_layer(in_channels=in_channels, num_blocks=num_blocks[1], stride=2)
- )
- in_channels *= 2 # 256
- self.dark4 = nn.SequentialCell(
- *self.make_group_layer(in_channels=in_channels, num_blocks=num_blocks[2], stride=2)
- )
- in_channels *= 2 # 512
- self.dark5 = nn.SequentialCell(
- *self.make_group_layer(in_channels=in_channels, num_blocks=num_blocks[3], stride=2),
- *self.make_spp_block([in_channels, in_channels * 2], in_channels * 2),
- )
-
- def make_group_layer(self, in_channels: int, num_blocks: int, stride: int = 1):
- "starts with conv layer then has `num_blocks` `ResLayer`"
- return [
- BaseConv(in_channels, in_channels * 2, ksize=3, stride=stride, act="lrelu"),
- *[(ResLayer(in_channels * 2)) for _ in range(num_blocks)],
- ]
-
- def make_spp_block(self, filters_list, in_filters):
- """ spatial pyramid pooling block"""
- m = nn.SequentialCell(
- *[
- BaseConv(in_filters, filters_list[0], 1, stride=1, act="lrelu"),
- BaseConv(filters_list[0], filters_list[1], 3, stride=1, act="lrelu"),
- SPPBottleneck(
- in_channels=filters_list[1],
- out_channels=filters_list[0],
- activation="lrelu",
- ),
- BaseConv(filters_list[0], filters_list[1], 3, stride=1, act="lrelu"),
- BaseConv(filters_list[1], filters_list[0], 1, stride=1, act="lrelu"),
- ]
- )
- return m
-
- def construct(self, x):
- """ forward """
- outputs = {}
- x = self.stem(x)
- outputs["stem"] = x
- x = self.dark2(x)
- outputs["dark2"] = x
- x = self.dark3(x)
- outputs["dark3"] = x
- x = self.dark4(x)
- outputs["dark4"] = x
- x = self.dark5(x)
- outputs["dark5"] = x
- return outputs["dark3"], outputs["dark4"], outputs["dark5"]
-
-
- class CSPDarknet(nn.Cell):
- """ Darknet with CSP block for yolox-s m l x"""
-
- def __init__(
- self,
- dep_mul,
- wid_mul,
- out_features=("dark3", "dark4", "dark5"),
- depthwise=False,
- act="silu"
- ):
- super(CSPDarknet, self).__init__()
- assert out_features, "please provide output features of Darknet"
- self.out_features = out_features
- Conv = DWConv if depthwise else BaseConv
- base_channels = int(wid_mul * 64)
- base_depth = max(round(dep_mul * 3), 1)
-
- # stem
- self.stem = Focus(3, base_channels, ksize=3, act=act)
-
- # dark2
- self.dark2 = nn.SequentialCell(
- Conv(base_channels, base_channels * 2, 3, 2, act=act),
- CSPLayer(
- base_channels * 2,
- base_channels * 2,
- n=base_depth,
- depthwise=depthwise,
- act=act,
- ),
- )
-
- # dark3
- self.dark3 = nn.SequentialCell(
- Conv(base_channels * 2, base_channels * 4, 3, 2, act=act),
- CSPLayer(
- base_channels * 4,
- base_channels * 4,
- n=base_depth * 3,
- depthwise=depthwise,
- act=act,
- ),
- )
-
- # dark4
- self.dark4 = nn.SequentialCell(
- Conv(base_channels * 4, base_channels * 8, 3, 2, act=act),
- CSPLayer(
- base_channels * 8,
- base_channels * 8,
- n=base_depth * 3,
- depthwise=depthwise,
- act=act,
- ),
- )
-
- # dark5
- self.dark5 = nn.SequentialCell(
- Conv(base_channels * 8, base_channels * 16, 3, 2, act=act),
- SPPBottleneck(base_channels * 16, base_channels * 16, activation=act),
- CSPLayer(
- base_channels * 16,
- base_channels * 16,
- n=base_depth,
- shortcut=False,
- depthwise=depthwise,
- act=act,
- ),
- )
-
- def construct(self, x):
- """ forward """
- outputs = {}
- x = self.stem(x)
- outputs["stem"] = x
- x = self.dark2(x)
- outputs["dark2"] = x
- x = self.dark3(x)
- outputs["dark3"] = x
- x = self.dark4(x)
- outputs["dark4"] = x
- x = self.dark5(x)
- outputs["dark5"] = x
- return outputs["dark3"], outputs["dark4"], outputs["dark5"]
|