darknet.py 6.8 KB
Newer Older
L
LielinJiang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
#  Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
#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.

15
import paddle
L
LielinJiang 已提交
16 17 18 19
import paddle.fluid as fluid
from paddle.fluid.param_attr import ParamAttr
from paddle.fluid.regularizer import L2Decay

20
from paddle.static import InputSpec
D
dengkaipeng 已提交
21
from paddle.fluid.dygraph.nn import Conv2D, BatchNorm
22
from paddle.utils.download import get_weights_path_from_url
L
LielinJiang 已提交
23 24 25 26

__all__ = ['DarkNet', 'darknet53']

# {num_layers: (url, md5)}
D
dengkaipeng 已提交
27
pretrain_infos = {
28 29
    53: ('https://paddlemodels.bj.bcebos.com/hapi/darknet53.pdparams',
         '2506357a5c31e865785112fc614a487d')
L
LielinJiang 已提交
30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72
}


class ConvBNLayer(fluid.dygraph.Layer):
    def __init__(self,
                 ch_in,
                 ch_out,
                 filter_size=3,
                 stride=1,
                 groups=1,
                 padding=0,
                 act="leaky"):
        super(ConvBNLayer, self).__init__()

        self.conv = Conv2D(
            num_channels=ch_in,
            num_filters=ch_out,
            filter_size=filter_size,
            stride=stride,
            padding=padding,
            groups=groups,
            param_attr=ParamAttr(
                initializer=fluid.initializer.Normal(0., 0.02)),
            bias_attr=False,
            act=None)
        self.batch_norm = BatchNorm(
            num_channels=ch_out,
            param_attr=ParamAttr(
                initializer=fluid.initializer.Normal(0., 0.02),
                regularizer=L2Decay(0.)),
            bias_attr=ParamAttr(
                initializer=fluid.initializer.Constant(0.0),
                regularizer=L2Decay(0.)))

        self.act = act

    def forward(self, inputs):
        out = self.conv(inputs)
        out = self.batch_norm(out)
        if self.act == 'leaky':
            out = fluid.layers.leaky_relu(x=out, alpha=0.1)
        return out

73

L
LielinJiang 已提交
74
class DownSample(fluid.dygraph.Layer):
75
    def __init__(self, ch_in, ch_out, filter_size=3, stride=2, padding=1):
L
LielinJiang 已提交
76 77 78 79 80 81 82 83 84 85

        super(DownSample, self).__init__()

        self.conv_bn_layer = ConvBNLayer(
            ch_in=ch_in,
            ch_out=ch_out,
            filter_size=filter_size,
            stride=stride,
            padding=padding)
        self.ch_out = ch_out
86

L
LielinJiang 已提交
87 88 89 90
    def forward(self, inputs):
        out = self.conv_bn_layer(inputs)
        return out

91

L
LielinJiang 已提交
92 93 94 95 96
class BasicBlock(fluid.dygraph.Layer):
    def __init__(self, ch_in, ch_out):
        super(BasicBlock, self).__init__()

        self.conv1 = ConvBNLayer(
97
            ch_in=ch_in, ch_out=ch_out, filter_size=1, stride=1, padding=0)
L
LielinJiang 已提交
98 99
        self.conv2 = ConvBNLayer(
            ch_in=ch_out,
100
            ch_out=ch_out * 2,
L
LielinJiang 已提交
101 102 103
            filter_size=3,
            stride=1,
            padding=1)
104

L
LielinJiang 已提交
105 106 107 108 109 110
    def forward(self, inputs):
        conv1 = self.conv1(inputs)
        conv2 = self.conv2(conv1)
        out = fluid.layers.elementwise_add(x=inputs, y=conv2, act=None)
        return out

111

L
LielinJiang 已提交
112 113
class LayerWarp(fluid.dygraph.Layer):
    def __init__(self, ch_in, ch_out, count):
114
        super(LayerWarp, self).__init__()
L
LielinJiang 已提交
115 116 117

        self.basicblock0 = BasicBlock(ch_in, ch_out)
        self.res_out_list = []
118
        for i in range(1, count):
L
LielinJiang 已提交
119
            res_out = self.add_sublayer("basic_block_%d" % (i),
120
                                        BasicBlock(ch_out * 2, ch_out))
L
LielinJiang 已提交
121 122
            self.res_out_list.append(res_out)
        self.ch_out = ch_out
123 124

    def forward(self, inputs):
L
LielinJiang 已提交
125 126 127 128 129 130 131 132 133
        y = self.basicblock0(inputs)
        for basic_block_i in self.res_out_list:
            y = basic_block_i(y)
        return y


DarkNet_cfg = {53: ([1, 2, 8, 8, 4])}


134
class DarkNet(fluid.dygraph.Layer):
L
LielinJiang 已提交
135 136 137 138 139
    """DarkNet model from
    `"YOLOv3: An Incremental Improvement" <https://arxiv.org/abs/1804.02767>`_

    Args:
        num_layers (int): layer number of DarkNet, only 53 supported currently, default: 53.
D
dengkaipeng 已提交
140
        ch_in (int): channel number of input data, default 3.
L
LielinJiang 已提交
141 142
    """

D
dengkaipeng 已提交
143
    def __init__(self, num_layers=53, ch_in=3):
L
LielinJiang 已提交
144 145 146 147 148 149
        super(DarkNet, self).__init__()
        assert num_layers in DarkNet_cfg.keys(), \
            "only support num_layers in {} currently" \
            .format(DarkNet_cfg.keys())
        self.stages = DarkNet_cfg[num_layers]
        self.stages = self.stages[0:5]
D
dengkaipeng 已提交
150

L
LielinJiang 已提交
151
        self.conv0 = ConvBNLayer(
152
            ch_in=ch_in, ch_out=32, filter_size=3, stride=1, padding=1)
L
LielinJiang 已提交
153

154
        self.downsample0 = DownSample(ch_in=32, ch_out=32 * 2)
L
LielinJiang 已提交
155 156
        self.darknet53_conv_block_list = []
        self.downsample_list = []
157
        ch_in = [64, 128, 256, 512, 1024]
L
LielinJiang 已提交
158
        for i, stage in enumerate(self.stages):
159 160 161 162
            conv_block = self.add_sublayer("stage_%d" % (i),
                                           LayerWarp(
                                               int(ch_in[i]), 32 * (2**i),
                                               stage))
L
LielinJiang 已提交
163 164 165 166 167
            self.darknet53_conv_block_list.append(conv_block)
        for i in range(len(self.stages) - 1):
            downsample = self.add_sublayer(
                "stage_%d_downsample" % i,
                DownSample(
168
                    ch_in=32 * (2**(i + 1)), ch_out=32 * (2**(i + 2))))
L
LielinJiang 已提交
169 170
            self.downsample_list.append(downsample)

171 172
    def forward(self, inputs):

L
LielinJiang 已提交
173 174
        out = self.conv0(inputs)
        out = self.downsample0(out)
D
dengkaipeng 已提交
175
        blocks = []
L
LielinJiang 已提交
176 177
        for i, conv_block_i in enumerate(self.darknet53_conv_block_list):
            out = conv_block_i(out)
D
dengkaipeng 已提交
178
            blocks.append(out)
L
LielinJiang 已提交
179 180
            if i < len(self.stages) - 1:
                out = self.downsample_list[i](out)
D
dengkaipeng 已提交
181
        return blocks[-1:-4:-1]
L
LielinJiang 已提交
182 183


D
dengkaipeng 已提交
184 185
def _darknet(num_layers=53, input_channels=3, pretrained=True):
    model = DarkNet(num_layers, input_channels)
L
LielinJiang 已提交
186
    if pretrained:
D
dengkaipeng 已提交
187 188 189 190
        assert num_layers in pretrain_infos.keys(), \
                "DarkNet{} do not have pretrained weights now, " \
                "pretrained should be set as False".format(num_layers)
        weight_path = get_weights_path_from_url(*(pretrain_infos[num_layers]))
L
LielinJiang 已提交
191 192
        assert weight_path.endswith('.pdparams'), \
                "suffix of weight must be .pdparams"
193 194
        weight_dict, _ = fluid.load_dygraph(weight_path[:-9])
        model.set_dict(weight_dict)
L
LielinJiang 已提交
195 196 197
    return model


D
dengkaipeng 已提交
198
def darknet53(input_channels=3, pretrained=True):
L
LielinJiang 已提交
199 200 201 202 203 204 205
    """DarkNet 53-layer model
    
    Args:
        input_channels (bool): channel number of input data, default 3. 
        pretrained (bool): If True, returns a model pre-trained on ImageNet,
            default True.
    """
D
dengkaipeng 已提交
206
    return _darknet(53, input_channels, pretrained)