sparse_encoder.py 7.91 KB
Newer Older
zhangwenwei's avatar
zhangwenwei committed
1
from torch import nn as nn
zhangwenwei's avatar
zhangwenwei committed
2

3
from mmdet3d.ops import SparseBasicBlock, make_sparse_convmodule
zhangwenwei's avatar
zhangwenwei committed
4
from mmdet3d.ops import spconv as spconv
zhangwenwei's avatar
zhangwenwei committed
5
6
7
from ..registry import MIDDLE_ENCODERS


8
@MIDDLE_ENCODERS.register_module()
zhangwenwei's avatar
zhangwenwei committed
9
class SparseEncoder(nn.Module):
zhangwenwei's avatar
zhangwenwei committed
10
    r"""Sparse encoder for SECOND and Part-A2.
wuyuefeng's avatar
wuyuefeng committed
11
12

    Args:
wangtai's avatar
wangtai committed
13
14
        in_channels (int): The number of input channels.
        sparse_shape (list[int]): The sparse shape of input tensor.
15
16
17
18
        order (list[str]): Order of conv module. Defaults to ('conv',
            'norm', 'act').
        norm_cfg (dict): Config of normalization layer. Defaults to
            dict(type='BN1d', eps=1e-3, momentum=0.01).
wangtai's avatar
wangtai committed
19
        base_channels (int): Out channels for conv_input layer.
20
            Defaults to 16.
wangtai's avatar
wangtai committed
21
        output_channels (int): Out channels for conv_out layer.
22
            Defaults to 128.
wuyuefeng's avatar
wuyuefeng committed
23
        encoder_channels (tuple[tuple[int]]):
wangtai's avatar
wangtai committed
24
25
            Convolutional channels of each encode block.
        encoder_paddings (tuple[tuple[int]]): Paddings of each encode block.
26
27
            Defaults to ((16, ), (32, 32, 32), (64, 64, 64), (64, 64, 64)).
        block_type (str): Type of the block to use. Defaults to 'conv_module'.
wuyuefeng's avatar
wuyuefeng committed
28
    """
zhangwenwei's avatar
zhangwenwei committed
29
30
31

    def __init__(self,
                 in_channels,
wuyuefeng's avatar
wuyuefeng committed
32
33
34
35
36
37
38
39
                 sparse_shape,
                 order=('conv', 'norm', 'act'),
                 norm_cfg=dict(type='BN1d', eps=1e-3, momentum=0.01),
                 base_channels=16,
                 output_channels=128,
                 encoder_channels=((16, ), (32, 32, 32), (64, 64, 64), (64, 64,
                                                                        64)),
                 encoder_paddings=((1, ), (1, 1, 1), (1, 1, 1), ((0, 1, 1), 1,
40
41
                                                                 1)),
                 block_type='conv_module'):
zhangwenwei's avatar
zhangwenwei committed
42
        super().__init__()
43
        assert block_type in ['conv_module', 'basicblock']
wuyuefeng's avatar
wuyuefeng committed
44
        self.sparse_shape = sparse_shape
zhangwenwei's avatar
zhangwenwei committed
45
        self.in_channels = in_channels
wuyuefeng's avatar
wuyuefeng committed
46
47
48
49
50
51
        self.order = order
        self.base_channels = base_channels
        self.output_channels = output_channels
        self.encoder_channels = encoder_channels
        self.encoder_paddings = encoder_paddings
        self.stage_num = len(self.encoder_channels)
zhangwenwei's avatar
zhangwenwei committed
52
        # Spconv init all weight on its own
wuyuefeng's avatar
wuyuefeng committed
53
54
55
56
57
58
59
60

        assert isinstance(order, tuple) and len(order) == 3
        assert set(order) == {'conv', 'norm', 'act'}

        if self.order[0] != 'conv':  # pre activate
            self.conv_input = make_sparse_convmodule(
                in_channels,
                self.base_channels,
zhangwenwei's avatar
zhangwenwei committed
61
62
63
                3,
                norm_cfg=norm_cfg,
                padding=1,
wuyuefeng's avatar
wuyuefeng committed
64
65
66
67
68
69
70
                indice_key='subm1',
                conv_type='SubMConv3d',
                order=('conv', ))
        else:  # post activate
            self.conv_input = make_sparse_convmodule(
                in_channels,
                self.base_channels,
zhangwenwei's avatar
zhangwenwei committed
71
72
73
                3,
                norm_cfg=norm_cfg,
                padding=1,
wuyuefeng's avatar
wuyuefeng committed
74
75
76
77
                indice_key='subm1',
                conv_type='SubMConv3d')

        encoder_out_channels = self.make_encoder_layers(
78
79
80
81
            make_sparse_convmodule,
            norm_cfg,
            self.base_channels,
            block_type=block_type)
wuyuefeng's avatar
wuyuefeng committed
82
83
84
85
86
87
88
89
90
91

        self.conv_out = make_sparse_convmodule(
            encoder_out_channels,
            self.output_channels,
            kernel_size=(3, 1, 1),
            stride=(2, 1, 1),
            norm_cfg=norm_cfg,
            padding=0,
            indice_key='spconv_down2',
            conv_type='SparseConv3d')
zhangwenwei's avatar
zhangwenwei committed
92
93

    def forward(self, voxel_features, coors, batch_size):
zhangwenwei's avatar
zhangwenwei committed
94
        """Forward of SparseEncoder.
wuyuefeng's avatar
wuyuefeng committed
95
96

        Args:
wangtai's avatar
wangtai committed
97
98
99
100
            voxel_features (torch.float32): Voxel features in shape (N, C).
            coors (torch.int32): Coordinates in shape (N, 4), \
                the columns in the order of (batch_idx, z_idx, y_idx, x_idx).
            batch_size (int): Batch size.
wuyuefeng's avatar
wuyuefeng committed
101
102

        Returns:
wangtai's avatar
wangtai committed
103
            dict: Backbone features.
zhangwenwei's avatar
zhangwenwei committed
104
105
106
107
108
109
110
        """
        coors = coors.int()
        input_sp_tensor = spconv.SparseConvTensor(voxel_features, coors,
                                                  self.sparse_shape,
                                                  batch_size)
        x = self.conv_input(input_sp_tensor)

wuyuefeng's avatar
wuyuefeng committed
111
112
113
114
        encode_features = []
        for encoder_layer in self.encoder_layers:
            x = encoder_layer(x)
            encode_features.append(x)
zhangwenwei's avatar
zhangwenwei committed
115
116
117

        # for detection head
        # [200, 176, 5] -> [200, 176, 2]
wuyuefeng's avatar
wuyuefeng committed
118
        out = self.conv_out(encode_features[-1])
zhangwenwei's avatar
zhangwenwei committed
119
120
121
122
123
124
125
        spatial_features = out.dense()

        N, C, D, H, W = spatial_features.shape
        spatial_features = spatial_features.view(N, C * D, H, W)

        return spatial_features

126
127
128
129
130
131
    def make_encoder_layers(self,
                            make_block,
                            norm_cfg,
                            in_channels,
                            block_type='conv_module',
                            conv_cfg=dict(type='SubMConv3d')):
zhangwenwei's avatar
zhangwenwei committed
132
        """make encoder layers using sparse convs.
wuyuefeng's avatar
wuyuefeng committed
133
134

        Args:
wangtai's avatar
wangtai committed
135
136
137
            make_block (method): A bounded function to build blocks.
            norm_cfg (dict[str]): Config of normalization layer.
            in_channels (int): The number of encoder input channels.
138
139
140
141
            block_type (str): Type of the block to use. Defaults to
                'conv_module'.
            conv_cfg (dict): Config of conv layer. Defaults to
                dict(type='SubMConv3d').
wuyuefeng's avatar
wuyuefeng committed
142
143

        Returns:
wangtai's avatar
wangtai committed
144
            int: The number of encoder output channels.
wuyuefeng's avatar
wuyuefeng committed
145
        """
146
        assert block_type in ['conv_module', 'basicblock']
wuyuefeng's avatar
wuyuefeng committed
147
148
149
150
151
152
153
154
        self.encoder_layers = spconv.SparseSequential()

        for i, blocks in enumerate(self.encoder_channels):
            blocks_list = []
            for j, out_channels in enumerate(tuple(blocks)):
                padding = tuple(self.encoder_paddings[i])[j]
                # each stage started with a spconv layer
                # except the first stage
155
                if i != 0 and j == 0 and block_type == 'conv_module':
wuyuefeng's avatar
wuyuefeng committed
156
157
158
159
160
161
162
163
164
165
                    blocks_list.append(
                        make_block(
                            in_channels,
                            out_channels,
                            3,
                            norm_cfg=norm_cfg,
                            stride=2,
                            padding=padding,
                            indice_key=f'spconv{i + 1}',
                            conv_type='SparseConv3d'))
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
                elif block_type == 'basicblock':
                    if j == len(blocks) - 1 and i != len(
                            self.encoder_channels) - 1:
                        blocks_list.append(
                            make_block(
                                in_channels,
                                out_channels,
                                3,
                                norm_cfg=norm_cfg,
                                stride=2,
                                padding=padding,
                                indice_key=f'spconv{i + 1}',
                                conv_type='SparseConv3d'))
                    else:
                        blocks_list.append(
                            SparseBasicBlock(
                                out_channels,
                                out_channels,
                                norm_cfg=norm_cfg,
                                conv_cfg=conv_cfg))
wuyuefeng's avatar
wuyuefeng committed
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
                else:
                    blocks_list.append(
                        make_block(
                            in_channels,
                            out_channels,
                            3,
                            norm_cfg=norm_cfg,
                            padding=padding,
                            indice_key=f'subm{i + 1}',
                            conv_type='SubMConv3d'))
                in_channels = out_channels
            stage_name = f'encoder_layer{i + 1}'
            stage_layers = spconv.SparseSequential(*blocks_list)
            self.encoder_layers.add_module(stage_name, stage_layers)
        return out_channels