| 8 | |
| 9 | |
| 10 | class Darknet(nn.Module): |
| 11 | # number of blocks from dark2 to dark5. |
| 12 | depth2blocks = {21: [1, 2, 2, 1], 53: [2, 8, 8, 4]} |
| 13 | |
| 14 | def __init__( |
| 15 | self, |
| 16 | depth, |
| 17 | in_channels=3, |
| 18 | stem_out_channels=32, |
| 19 | out_features=("dark3", "dark4", "dark5"), |
| 20 | ): |
| 21 | """ |
| 22 | Args: |
| 23 | depth (int): depth of darknet used in model, usually use [21, 53] for this param. |
| 24 | in_channels (int): number of input channels, for example, use 3 for RGB image. |
| 25 | stem_out_channels (int): number of output chanels of darknet stem. |
| 26 | It decides channels of darknet layer2 to layer5. |
| 27 | out_features (Tuple[str]): desired output layer name. |
| 28 | """ |
| 29 | super().__init__() |
| 30 | assert out_features, "please provide output features of Darknet" |
| 31 | self.out_features = out_features |
| 32 | self.stem = nn.Sequential( |
| 33 | BaseConv(in_channels, stem_out_channels, ksize=3, stride=1, act="lrelu"), |
| 34 | *self.make_group_layer(stem_out_channels, num_blocks=1, stride=2), |
| 35 | ) |
| 36 | in_channels = stem_out_channels * 2 # 64 |
| 37 | |
| 38 | num_blocks = Darknet.depth2blocks[depth] |
| 39 | # create darknet with `stem_out_channels` and `num_blocks` layers. |
| 40 | # to make model structure more clear, we don't use `for` statement in python. |
| 41 | self.dark2 = nn.Sequential( |
| 42 | *self.make_group_layer(in_channels, num_blocks[0], stride=2) |
| 43 | ) |
| 44 | in_channels *= 2 # 128 |
| 45 | self.dark3 = nn.Sequential( |
| 46 | *self.make_group_layer(in_channels, num_blocks[1], stride=2) |
| 47 | ) |
| 48 | in_channels *= 2 # 256 |
| 49 | self.dark4 = nn.Sequential( |
| 50 | *self.make_group_layer(in_channels, num_blocks[2], stride=2) |
| 51 | ) |
| 52 | in_channels *= 2 # 512 |
| 53 | |
| 54 | self.dark5 = nn.Sequential( |
| 55 | *self.make_group_layer(in_channels, num_blocks[3], stride=2), |
| 56 | *self.make_spp_block([in_channels, in_channels * 2], in_channels * 2), |
| 57 | ) |
| 58 | |
| 59 | def make_group_layer(self, in_channels: int, num_blocks: int, stride: int = 1): |
| 60 | "starts with conv layer then has `num_blocks` `ResLayer`" |
| 61 | return [ |
| 62 | BaseConv(in_channels, in_channels * 2, ksize=3, stride=stride, act="lrelu"), |
| 63 | *[(ResLayer(in_channels * 2)) for _ in range(num_blocks)], |
| 64 | ] |
| 65 | |
| 66 | def make_spp_block(self, filters_list, in_filters): |
| 67 | m = nn.Sequential( |