second_fpn.py 4.71 KB
Newer Older
zhangwenwei's avatar
zhangwenwei committed
1
2
import torch
import torch.nn as nn
zhangwenwei's avatar
zhangwenwei committed
3
4
from mmcv.cnn import (build_norm_layer, build_upsample_layer, constant_init,
                      is_norm, kaiming_init)
zhangwenwei's avatar
zhangwenwei committed
5

zhangwenwei's avatar
zhangwenwei committed
6
from mmdet.models import NECKS
zhangwenwei's avatar
zhangwenwei committed
7
8
9
from .. import builder


10
@NECKS.register_module()
zhangwenwei's avatar
zhangwenwei committed
11
class SECONDFPN(nn.Module):
zhangwenwei's avatar
zhangwenwei committed
12
    """FPN used in SECOND/PointPillars/PartA2/MVXNet
zhangwenwei's avatar
zhangwenwei committed
13
14
15
16
17
18
19

    Args:
        in_channels (list[int]): Input channels of multi-scale feature maps
        out_channels (list[int]): Output channels of feature maps
        upsample_strides (list[int]): Strides used to upsample the feature maps
        norm_cfg (dict): Config dict of normalization layers
        upsample_cfg (dict): Config dict of upsample layers
zhangwenwei's avatar
zhangwenwei committed
20
21
22
23
    """

    def __init__(self,
                 in_channels=[128, 128, 256],
zhangwenwei's avatar
zhangwenwei committed
24
                 out_channels=[256, 256, 256],
zhangwenwei's avatar
zhangwenwei committed
25
                 upsample_strides=[1, 2, 4],
zhangwenwei's avatar
zhangwenwei committed
26
27
                 norm_cfg=dict(type='BN', eps=1e-3, momentum=0.01),
                 upsample_cfg=dict(type='deconv', bias=False)):
zhangwenwei's avatar
zhangwenwei committed
28
29
30
        # if for GroupNorm,
        # cfg is dict(type='GN', num_groups=num_groups, eps=1e-3, affine=True)
        super(SECONDFPN, self).__init__()
zhangwenwei's avatar
zhangwenwei committed
31
        assert len(out_channels) == len(upsample_strides) == len(in_channels)
zhangwenwei's avatar
zhangwenwei committed
32
        self.in_channels = in_channels
zhangwenwei's avatar
zhangwenwei committed
33
        self.out_channels = out_channels
zhangwenwei's avatar
zhangwenwei committed
34
35

        deblocks = []
zhangwenwei's avatar
zhangwenwei committed
36
        for i, out_channel in enumerate(out_channels):
zhangwenwei's avatar
zhangwenwei committed
37
38
            upsample_layer = build_upsample_layer(
                upsample_cfg,
zhangwenwei's avatar
zhangwenwei committed
39
40
                in_channels=in_channels[i],
                out_channels=out_channel,
zhangwenwei's avatar
zhangwenwei committed
41
                kernel_size=upsample_strides[i],
zhangwenwei's avatar
zhangwenwei committed
42
                stride=upsample_strides[i])
zhangwenwei's avatar
zhangwenwei committed
43
44
45
            deblock = nn.Sequential(upsample_layer,
                                    build_norm_layer(norm_cfg, out_channel)[1],
                                    nn.ReLU(inplace=True))
zhangwenwei's avatar
zhangwenwei committed
46
47
48
            deblocks.append(deblock)
        self.deblocks = nn.ModuleList(deblocks)

zhangwenwei's avatar
zhangwenwei committed
49
50
51
52
    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                kaiming_init(m)
zhangwenwei's avatar
zhangwenwei committed
53
            elif is_norm(m):
zhangwenwei's avatar
zhangwenwei committed
54
                constant_init(m, 1)
zhangwenwei's avatar
zhangwenwei committed
55

zhangwenwei's avatar
zhangwenwei committed
56
57
58
    def forward(self, x):
        assert len(x) == len(self.in_channels)
        ups = [deblock(x[i]) for i, deblock in enumerate(self.deblocks)]
zhangwenwei's avatar
zhangwenwei committed
59
60

        if len(ups) > 1:
zhangwenwei's avatar
zhangwenwei committed
61
            out = torch.cat(ups, dim=1)
zhangwenwei's avatar
zhangwenwei committed
62
        else:
zhangwenwei's avatar
zhangwenwei committed
63
64
            out = ups[0]
        return [out]
zhangwenwei's avatar
zhangwenwei committed
65
66


67
@NECKS.register_module()
zhangwenwei's avatar
zhangwenwei committed
68
class SECONDFusionFPN(SECONDFPN):
zhangwenwei's avatar
zhangwenwei committed
69
70
71
72
73
74
75
76
77
78
79
    """FPN used in multi-modality SECOND/PointPillars

    Args:
        in_channels (list[int]): Input channels of multi-scale feature maps
        out_channels (list[int]): Output channels of feature maps
        upsample_strides (list[int]): Strides used to upsample the feature maps
        norm_cfg (dict): Config dict of normalization layers
        upsample_cfg (dict): Config dict of upsample layers
        downsample_rates (list[int]): The downsample rate of feature map in
            comparison to the original voxelization input
        fusion_layer (dict): Config dict of fusion layers
zhangwenwei's avatar
zhangwenwei committed
80
81
82
83
    """

    def __init__(self,
                 in_channels=[128, 128, 256],
zhangwenwei's avatar
zhangwenwei committed
84
                 out_channels=[256, 256, 256],
zhangwenwei's avatar
zhangwenwei committed
85
86
                 upsample_strides=[1, 2, 4],
                 norm_cfg=dict(type='BN', eps=1e-3, momentum=0.01),
zhangwenwei's avatar
zhangwenwei committed
87
88
89
90
91
92
                 upsample_cfg=dict(type='deconv', bias=False),
                 downsample_rates=[40, 8, 8],
                 fusion_layer=None):
        super(SECONDFusionFPN,
              self).__init__(in_channels, out_channels, upsample_strides,
                             norm_cfg, upsample_cfg)
zhangwenwei's avatar
zhangwenwei committed
93
94
95
        self.fusion_layer = None
        if fusion_layer is not None:
            self.fusion_layer = builder.build_fusion_layer(fusion_layer)
zhangwenwei's avatar
zhangwenwei committed
96
        self.downsample_rates = downsample_rates
zhangwenwei's avatar
zhangwenwei committed
97
98

    def forward(self,
zhangwenwei's avatar
zhangwenwei committed
99
                x,
zhangwenwei's avatar
zhangwenwei committed
100
101
102
                coors=None,
                points=None,
                img_feats=None,
zhangwenwei's avatar
zhangwenwei committed
103
                img_metas=None):
zhangwenwei's avatar
zhangwenwei committed
104
105
        assert len(x) == len(self.in_channels)
        ups = [deblock(x[i]) for i, deblock in enumerate(self.deblocks)]
zhangwenwei's avatar
zhangwenwei committed
106
107

        if len(ups) > 1:
zhangwenwei's avatar
zhangwenwei committed
108
            out = torch.cat(ups, dim=1)
zhangwenwei's avatar
zhangwenwei committed
109
        else:
zhangwenwei's avatar
zhangwenwei committed
110
            out = ups[0]
zhangwenwei's avatar
zhangwenwei committed
111
112
113
114
        if (self.fusion_layer is not None and img_feats is not None):
            downsample_pts_coors = torch.zeros_like(coors)
            downsample_pts_coors[:, 0] = coors[:, 0]
            downsample_pts_coors[:, 1] = (
zhangwenwei's avatar
zhangwenwei committed
115
                coors[:, 1] / self.downsample_rates[0])
zhangwenwei's avatar
zhangwenwei committed
116
            downsample_pts_coors[:, 2] = (
zhangwenwei's avatar
zhangwenwei committed
117
                coors[:, 2] / self.downsample_rates[1])
zhangwenwei's avatar
zhangwenwei committed
118
            downsample_pts_coors[:, 3] = (
zhangwenwei's avatar
zhangwenwei committed
119
                coors[:, 3] / self.downsample_rates[2])
zhangwenwei's avatar
zhangwenwei committed
120
            # fusion for each point
zhangwenwei's avatar
zhangwenwei committed
121
            out = self.fusion_layer(img_feats, points, out,
zhangwenwei's avatar
zhangwenwei committed
122
                                    downsample_pts_coors, img_metas)
zhangwenwei's avatar
zhangwenwei committed
123
        return [out]