sparse_unet.py 11.2 KB
Newer Older
wuyuefeng's avatar
wuyuefeng committed
1
import torch
zhangwenwei's avatar
zhangwenwei committed
2
from torch import nn as nn
wuyuefeng's avatar
wuyuefeng committed
3

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


9
@MIDDLE_ENCODERS.register_module()
wuyuefeng's avatar
wuyuefeng committed
10
class SparseUNet(nn.Module):
zhangwenwei's avatar
zhangwenwei committed
11
    r"""SparseUNet for PartA^2.
wuyuefeng's avatar
wuyuefeng committed
12

zhangwenwei's avatar
zhangwenwei committed
13
    See the `paper <https://arxiv.org/abs/1907.03670>`_ for more detials.
wuyuefeng's avatar
wuyuefeng committed
14
15

    Args:
wangtai's avatar
wangtai committed
16
17
18
19
20
        in_channels (int): The number of input channels.
        sparse_shape (list[int]): The sparse shape of input tensor.
        norm_cfg (dict): Config of normalization layer.
        base_channels (int): Out channels for conv_input layer.
        output_channels (int): Out channels for conv_out layer.
wuyuefeng's avatar
wuyuefeng committed
21
        encoder_channels (tuple[tuple[int]]):
wangtai's avatar
wangtai committed
22
23
            Convolutional channels of each encode block.
        encoder_paddings (tuple[tuple[int]]): Paddings of each encode block.
wuyuefeng's avatar
wuyuefeng committed
24
        decoder_channels (tuple[tuple[int]]):
wangtai's avatar
wangtai committed
25
26
            Convolutional channels of each decode block.
        decoder_paddings (tuple[tuple[int]]): Paddings of each decode block.
wuyuefeng's avatar
wuyuefeng committed
27
    """
wuyuefeng's avatar
wuyuefeng committed
28
29
30

    def __init__(self,
                 in_channels,
wuyuefeng's avatar
wuyuefeng committed
31
32
                 sparse_shape,
                 order=('conv', 'norm', 'act'),
wuyuefeng's avatar
wuyuefeng committed
33
34
                 norm_cfg=dict(type='BN1d', eps=1e-3, momentum=0.01),
                 base_channels=16,
35
36
37
38
39
40
41
42
                 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,
                                                                 1)),
                 decoder_channels=((64, 64, 64), (64, 64, 32), (32, 32, 16),
                                   (16, 16, 16)),
                 decoder_paddings=((1, 0), (1, 0), (0, 0), (0, 1))):
wuyuefeng's avatar
wuyuefeng committed
43
        super().__init__()
wuyuefeng's avatar
wuyuefeng committed
44
        self.sparse_shape = sparse_shape
wuyuefeng's avatar
wuyuefeng committed
45
        self.in_channels = in_channels
wuyuefeng's avatar
wuyuefeng committed
46
        self.order = order
wuyuefeng's avatar
wuyuefeng committed
47
        self.base_channels = base_channels
48
49
50
51
52
53
        self.output_channels = output_channels
        self.encoder_channels = encoder_channels
        self.encoder_paddings = encoder_paddings
        self.decoder_channels = decoder_channels
        self.decoder_paddings = decoder_paddings
        self.stage_num = len(self.encoder_channels)
wuyuefeng's avatar
wuyuefeng committed
54
55
        # Spconv init all weight on its own

wuyuefeng's avatar
wuyuefeng committed
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
        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,
                3,
                norm_cfg=norm_cfg,
                padding=1,
                indice_key='subm1',
                conv_type='SubMConv3d',
                order=('conv', ))
        else:  # post activate
            self.conv_input = make_sparse_convmodule(
                in_channels,
                self.base_channels,
                3,
                norm_cfg=norm_cfg,
                padding=1,
                indice_key='subm1',
                conv_type='SubMConv3d')
wuyuefeng's avatar
wuyuefeng committed
78

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

        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')
wuyuefeng's avatar
wuyuefeng committed
93
94

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

        Args:
zhangwenwei's avatar
zhangwenwei committed
98
99
100
101
            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
102
103

        Returns:
zhangwenwei's avatar
zhangwenwei committed
104
            dict[str, torch.Tensor]: Backbone features.
wuyuefeng's avatar
wuyuefeng committed
105
106
107
108
109
110
111
        """
        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
112
        encode_features = []
wuyuefeng's avatar
wuyuefeng committed
113
114
        for encoder_layer in self.encoder_layers:
            x = encoder_layer(x)
wuyuefeng's avatar
wuyuefeng committed
115
            encode_features.append(x)
wuyuefeng's avatar
wuyuefeng committed
116
117
118

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

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

wuyuefeng's avatar
wuyuefeng committed
125
        # for segmentation head, with output shape:
wuyuefeng's avatar
wuyuefeng committed
126
127
128
129
        # [400, 352, 11] <- [200, 176, 5]
        # [800, 704, 21] <- [400, 352, 11]
        # [1600, 1408, 41] <- [800, 704, 21]
        # [1600, 1408, 41] <- [1600, 1408, 41]
wuyuefeng's avatar
wuyuefeng committed
130
131
132
        decode_features = []
        x = encode_features[-1]
        for i in range(self.stage_num, 0, -1):
wuyuefeng's avatar
wuyuefeng committed
133
134
135
136
            x = self.decoder_layer_forward(encode_features[i - 1], x,
                                           getattr(self, f'lateral_layer{i}'),
                                           getattr(self, f'merge_layer{i}'),
                                           getattr(self, f'upsample_layer{i}'))
wuyuefeng's avatar
wuyuefeng committed
137
            decode_features.append(x)
wuyuefeng's avatar
wuyuefeng committed
138

wuyuefeng's avatar
wuyuefeng committed
139
        seg_features = decode_features[-1].features
wuyuefeng's avatar
wuyuefeng committed
140

wuyuefeng's avatar
wuyuefeng committed
141
142
        ret = dict(
            spatial_features=spatial_features, seg_features=seg_features)
wuyuefeng's avatar
wuyuefeng committed
143
144
145

        return ret

wuyuefeng's avatar
wuyuefeng committed
146
147
    def decoder_layer_forward(self, x_lateral, x_bottom, lateral_layer,
                              merge_layer, upsample_layer):
wuyuefeng's avatar
wuyuefeng committed
148
149
150
        """Forward of upsample and residual block.

        Args:
zhangwenwei's avatar
zhangwenwei committed
151
152
153
154
155
            x_lateral (:obj:`SparseConvTensor`): Lateral tensor.
            x_bottom (:obj:`SparseConvTensor`): Feature from bottom layer.
            lateral_layer (SparseBasicBlock): Convolution for lateral tensor.
            merge_layer (SparseSequential): Convolution for merging features.
            upsample_layer (SparseSequential): Convolution for upsampling.
wuyuefeng's avatar
wuyuefeng committed
156
157

        Returns:
zhangwenwei's avatar
zhangwenwei committed
158
            :obj:`SparseConvTensor`: Upsampled feature.
wuyuefeng's avatar
wuyuefeng committed
159
        """
wuyuefeng's avatar
wuyuefeng committed
160
161
162
163
164
        x = lateral_layer(x_lateral)
        x.features = torch.cat((x_bottom.features, x.features), dim=1)
        x_merge = merge_layer(x)
        x = self.reduce_channel(x, x_merge.features.shape[1])
        x.features = x_merge.features + x.features
wuyuefeng's avatar
wuyuefeng committed
165
        x = upsample_layer(x)
wuyuefeng's avatar
wuyuefeng committed
166
167
168
        return x

    @staticmethod
wuyuefeng's avatar
wuyuefeng committed
169
170
    def reduce_channel(x, out_channels):
        """reduce channel for element-wise addition.
wuyuefeng's avatar
wuyuefeng committed
171
172

        Args:
zhangwenwei's avatar
zhangwenwei committed
173
174
175
            x (:obj:`SparseConvTensor`): Sparse tensor, ``x.features``
                are in shape (N, C1).
            out_channels (int): The number of channel after reduction.
wuyuefeng's avatar
wuyuefeng committed
176
177

        Returns:
zhangwenwei's avatar
zhangwenwei committed
178
            :obj:`SparseConvTensor`: Channel reduced feature.
wuyuefeng's avatar
wuyuefeng committed
179
180
181
        """
        features = x.features
        n, in_channels = features.shape
wuyuefeng's avatar
wuyuefeng committed
182
183
        assert (in_channels % out_channels
                == 0) and (in_channels >= out_channels)
wuyuefeng's avatar
wuyuefeng committed
184
185
186
187

        x.features = features.view(n, out_channels, -1).sum(dim=2)
        return x

188
    def make_encoder_layers(self, make_block, norm_cfg, in_channels):
zhangwenwei's avatar
zhangwenwei committed
189
        """make encoder layers using sparse convs.
wuyuefeng's avatar
wuyuefeng committed
190
191

        Args:
zhangwenwei's avatar
zhangwenwei committed
192
193
194
            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.
wuyuefeng's avatar
wuyuefeng committed
195
196

        Returns:
wangtai's avatar
wangtai committed
197
            int: The number of encoder output channels.
wuyuefeng's avatar
wuyuefeng committed
198
        """
wuyuefeng's avatar
wuyuefeng committed
199
        self.encoder_layers = spconv.SparseSequential()
wuyuefeng's avatar
wuyuefeng committed
200

201
        for i, blocks in enumerate(self.encoder_channels):
wuyuefeng's avatar
wuyuefeng committed
202
203
            blocks_list = []
            for j, out_channels in enumerate(tuple(blocks)):
204
                padding = tuple(self.encoder_paddings[i])[j]
wuyuefeng's avatar
wuyuefeng committed
205
206
207
208
209
210
211
212
213
214
215
                # each stage started with a spconv layer
                # except the first stage
                if i != 0 and j == 0:
                    blocks_list.append(
                        make_block(
                            in_channels,
                            out_channels,
                            3,
                            norm_cfg=norm_cfg,
                            stride=2,
                            padding=padding,
216
                            indice_key=f'spconv{i + 1}',
wuyuefeng's avatar
wuyuefeng committed
217
                            conv_type='SparseConv3d'))
wuyuefeng's avatar
wuyuefeng committed
218
219
220
221
222
223
224
225
                else:
                    blocks_list.append(
                        make_block(
                            in_channels,
                            out_channels,
                            3,
                            norm_cfg=norm_cfg,
                            padding=padding,
wuyuefeng's avatar
wuyuefeng committed
226
227
                            indice_key=f'subm{i + 1}',
                            conv_type='SubMConv3d'))
wuyuefeng's avatar
wuyuefeng committed
228
                in_channels = out_channels
229
            stage_name = f'encoder_layer{i + 1}'
wuyuefeng's avatar
wuyuefeng committed
230
            stage_layers = spconv.SparseSequential(*blocks_list)
wuyuefeng's avatar
wuyuefeng committed
231
            self.encoder_layers.add_module(stage_name, stage_layers)
wuyuefeng's avatar
wuyuefeng committed
232
233
        return out_channels

234
    def make_decoder_layers(self, make_block, norm_cfg, in_channels):
zhangwenwei's avatar
zhangwenwei committed
235
        """make decoder layers using sparse convs.
wuyuefeng's avatar
wuyuefeng committed
236
237

        Args:
zhangwenwei's avatar
zhangwenwei committed
238
239
240
            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.
wuyuefeng's avatar
wuyuefeng committed
241
242

        Returns:
zhangwenwei's avatar
zhangwenwei committed
243
            int: The number of encoder output channels.
wuyuefeng's avatar
wuyuefeng committed
244
        """
245
246
247
        block_num = len(self.decoder_channels)
        for i, block_channels in enumerate(self.decoder_channels):
            paddings = self.decoder_paddings[i]
wuyuefeng's avatar
wuyuefeng committed
248
            setattr(
249
                self, f'lateral_layer{block_num - i}',
wuyuefeng's avatar
wuyuefeng committed
250
251
252
253
                SparseBasicBlock(
                    in_channels,
                    block_channels[0],
                    conv_cfg=dict(
254
                        type='SubMConv3d', indice_key=f'subm{block_num - i}'),
wuyuefeng's avatar
wuyuefeng committed
255
256
                    norm_cfg=norm_cfg))
            setattr(
257
                self, f'merge_layer{block_num - i}',
wuyuefeng's avatar
wuyuefeng committed
258
259
260
261
262
263
                make_block(
                    in_channels * 2,
                    block_channels[1],
                    3,
                    norm_cfg=norm_cfg,
                    padding=paddings[0],
wuyuefeng's avatar
wuyuefeng committed
264
265
                    indice_key=f'subm{block_num - i}',
                    conv_type='SubMConv3d'))
wuyuefeng's avatar
wuyuefeng committed
266
267
268
269
270
271
272
273
274
            if block_num - i != 1:
                setattr(
                    self, f'upsample_layer{block_num - i}',
                    make_block(
                        in_channels,
                        block_channels[2],
                        3,
                        norm_cfg=norm_cfg,
                        indice_key=f'spconv{block_num - i}',
wuyuefeng's avatar
wuyuefeng committed
275
                        conv_type='SparseInverseConv3d'))
wuyuefeng's avatar
wuyuefeng committed
276
277
            else:
                # use submanifold conv instead of inverse conv
wuyuefeng's avatar
wuyuefeng committed
278
                # in the last block
wuyuefeng's avatar
wuyuefeng committed
279
280
281
282
283
284
285
286
287
                setattr(
                    self, f'upsample_layer{block_num - i}',
                    make_block(
                        in_channels,
                        block_channels[2],
                        3,
                        norm_cfg=norm_cfg,
                        padding=paddings[1],
                        indice_key='subm1',
wuyuefeng's avatar
wuyuefeng committed
288
                        conv_type='SubMConv3d'))
wuyuefeng's avatar
wuyuefeng committed
289
            in_channels = block_channels[2]