second.py 3.12 KB
Newer Older
1
import warnings
zhangwenwei's avatar
zhangwenwei committed
2
from mmcv.cnn import build_conv_layer, build_norm_layer
3
from mmcv.runner import BaseModule
zhangwenwei's avatar
zhangwenwei committed
4
from torch import nn as nn
zhangwenwei's avatar
zhangwenwei committed
5

zhangwenwei's avatar
zhangwenwei committed
6
from mmdet.models import BACKBONES
zhangwenwei's avatar
zhangwenwei committed
7
8


9
@BACKBONES.register_module()
10
class SECOND(BaseModule):
zhangwenwei's avatar
zhangwenwei committed
11
    """Backbone network for SECOND/PointPillars/PartA2/MVXNet.
zhangwenwei's avatar
zhangwenwei committed
12
13

    Args:
wangtai's avatar
wangtai committed
14
15
16
17
18
19
        in_channels (int): Input channels.
        out_channels (list[int]): Output channels for multi-scale feature maps.
        layer_nums (list[int]): Number of layers in each stage.
        layer_strides (list[int]): Strides of each stage.
        norm_cfg (dict): Config dict of normalization layers.
        conv_cfg (dict): Config dict of convolutional layers.
zhangwenwei's avatar
zhangwenwei committed
20
21
22
23
    """

    def __init__(self,
                 in_channels=128,
zhangwenwei's avatar
zhangwenwei committed
24
                 out_channels=[128, 128, 256],
zhangwenwei's avatar
zhangwenwei committed
25
26
                 layer_nums=[3, 5, 5],
                 layer_strides=[2, 2, 2],
zhangwenwei's avatar
zhangwenwei committed
27
                 norm_cfg=dict(type='BN', eps=1e-3, momentum=0.01),
28
29
30
31
                 conv_cfg=dict(type='Conv2d', bias=False),
                 init_cfg=None,
                 pretrained=None):
        super(SECOND, self).__init__(init_cfg=init_cfg)
zhangwenwei's avatar
zhangwenwei committed
32
        assert len(layer_strides) == len(layer_nums)
zhangwenwei's avatar
zhangwenwei committed
33
        assert len(out_channels) == len(layer_nums)
zhangwenwei's avatar
zhangwenwei committed
34

zhangwenwei's avatar
zhangwenwei committed
35
        in_filters = [in_channels, *out_channels[:-1]]
zhangwenwei's avatar
zhangwenwei committed
36
37
38
39
40
        # note that when stride > 1, conv2d with same padding isn't
        # equal to pad-conv2d. we should use pad-conv2d.
        blocks = []
        for i, layer_num in enumerate(layer_nums):
            block = [
zhangwenwei's avatar
zhangwenwei committed
41
42
43
44
45
46
47
48
                build_conv_layer(
                    conv_cfg,
                    in_filters[i],
                    out_channels[i],
                    3,
                    stride=layer_strides[i],
                    padding=1),
                build_norm_layer(norm_cfg, out_channels[i])[1],
zhangwenwei's avatar
zhangwenwei committed
49
50
51
52
                nn.ReLU(inplace=True),
            ]
            for j in range(layer_num):
                block.append(
zhangwenwei's avatar
zhangwenwei committed
53
54
55
56
57
58
59
                    build_conv_layer(
                        conv_cfg,
                        out_channels[i],
                        out_channels[i],
                        3,
                        padding=1))
                block.append(build_norm_layer(norm_cfg, out_channels[i])[1])
zhangwenwei's avatar
zhangwenwei committed
60
61
62
63
64
65
66
                block.append(nn.ReLU(inplace=True))

            block = nn.Sequential(*block)
            blocks.append(block)

        self.blocks = nn.ModuleList(blocks)

67
68
        assert not (init_cfg and pretrained), \
            'init_cfg and pretrained cannot be setting at the same time'
zhangwenwei's avatar
zhangwenwei committed
69
        if isinstance(pretrained, str):
70
71
72
73
74
            warnings.warn('DeprecationWarning: pretrained is a deprecated, '
                          'please use "init_cfg" instead')
            self.init_cfg = dict(type='Pretrained', checkpoint=pretrained)
        else:
            self.init_cfg = dict(type='Kaiming', layer='Conv2d')
zhangwenwei's avatar
zhangwenwei committed
75
76

    def forward(self, x):
77
78
79
80
81
82
83
84
        """Forward function.

        Args:
            x (torch.Tensor): Input with shape (N, C, H, W).

        Returns:
            tuple[torch.Tensor]: Multi-scale features.
        """
zhangwenwei's avatar
zhangwenwei committed
85
86
87
88
89
        outs = []
        for i in range(len(self.blocks)):
            x = self.blocks[i](x)
            outs.append(x)
        return tuple(outs)