sparse_unet.py 11 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
    """SparseUNet for PartA^2.
wuyuefeng's avatar
wuyuefeng committed
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

    See https://arxiv.org/abs/1907.03670 for more detials.

    Args:
        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
        encoder_channels (tuple[tuple[int]]):
            conv channels of each encode block
        encoder_paddings (tuple[tuple[int]]): paddings of each encode block
        decoder_channels (tuple[tuple[int]]):
            conv channels of each decode block
        decoder_paddings (tuple[tuple[int]]): paddings of each decode block
    """
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
98
99
100
101
102
103
104
105
106
107
108
109
110

        Args:
            voxel_features (torch.float32): shape [N, C]
            coors (torch.int32): shape [N, 4](batch_idx, z_idx, y_idx, x_idx)
            batch_size (int): batch size

        Returns:
            dict: backbone features
        """
        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
        encode_features = []
wuyuefeng's avatar
wuyuefeng committed
112
113
        for encoder_layer in self.encoder_layers:
            x = encoder_layer(x)
wuyuefeng's avatar
wuyuefeng committed
114
            encode_features.append(x)
wuyuefeng's avatar
wuyuefeng 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])
wuyuefeng's avatar
wuyuefeng committed
119
120
121
122
123
        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
124
        # for segmentation head, with output shape:
wuyuefeng's avatar
wuyuefeng committed
125
126
127
128
        # [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
129
130
131
        decode_features = []
        x = encode_features[-1]
        for i in range(self.stage_num, 0, -1):
wuyuefeng's avatar
wuyuefeng committed
132
133
134
135
            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
136
            decode_features.append(x)
wuyuefeng's avatar
wuyuefeng committed
137

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

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

        return ret

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

        Args:
            x_lateral (SparseConvTensor): lateral tensor
wuyuefeng's avatar
wuyuefeng committed
151
            x_bottom (SparseConvTensor): feature from bottom layer
wuyuefeng's avatar
wuyuefeng committed
152
153
154
            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
155
156
157
158

        Returns:
            SparseConvTensor: upsampled feature
        """
wuyuefeng's avatar
wuyuefeng committed
159
160
161
162
163
        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
164
        x = upsample_layer(x)
wuyuefeng's avatar
wuyuefeng committed
165
166
167
        return x

    @staticmethod
wuyuefeng's avatar
wuyuefeng committed
168
169
    def reduce_channel(x, out_channels):
        """reduce channel for element-wise addition.
wuyuefeng's avatar
wuyuefeng committed
170
171
172
173
174
175
176
177
178
179

        Args:
            x (SparseConvTensor): x.features (N, C1)
            out_channels (int): the number of channel after reduction

        Returns:
            SparseConvTensor: channel reduced feature
        """
        features = x.features
        n, in_channels = features.shape
wuyuefeng's avatar
wuyuefeng committed
180
181
        assert (in_channels % out_channels
                == 0) and (in_channels >= out_channels)
wuyuefeng's avatar
wuyuefeng committed
182
183
184
185

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

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

        Args:
            make_block (method): a bounded function to build blocks
191
            norm_cfg (dict[str]): config of normalization layer
wuyuefeng's avatar
wuyuefeng committed
192
193
194
195
196
            in_channels (int): the number of encoder input channels

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

199
        for i, blocks in enumerate(self.encoder_channels):
wuyuefeng's avatar
wuyuefeng committed
200
201
            blocks_list = []
            for j, out_channels in enumerate(tuple(blocks)):
202
                padding = tuple(self.encoder_paddings[i])[j]
wuyuefeng's avatar
wuyuefeng committed
203
204
205
206
207
208
209
210
211
212
213
                # 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,
214
                            indice_key=f'spconv{i + 1}',
wuyuefeng's avatar
wuyuefeng committed
215
                            conv_type='SparseConv3d'))
wuyuefeng's avatar
wuyuefeng committed
216
217
218
219
220
221
222
223
                else:
                    blocks_list.append(
                        make_block(
                            in_channels,
                            out_channels,
                            3,
                            norm_cfg=norm_cfg,
                            padding=padding,
wuyuefeng's avatar
wuyuefeng committed
224
225
                            indice_key=f'subm{i + 1}',
                            conv_type='SubMConv3d'))
wuyuefeng's avatar
wuyuefeng committed
226
                in_channels = out_channels
227
            stage_name = f'encoder_layer{i + 1}'
wuyuefeng's avatar
wuyuefeng committed
228
            stage_layers = spconv.SparseSequential(*blocks_list)
wuyuefeng's avatar
wuyuefeng committed
229
            self.encoder_layers.add_module(stage_name, stage_layers)
wuyuefeng's avatar
wuyuefeng committed
230
231
        return out_channels

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

        Args:
            make_block (method): a bounded function to build blocks
237
            norm_cfg (dict[str]): config of normalization layer
wuyuefeng's avatar
wuyuefeng committed
238
239
240
241
242
            in_channels (int): the number of encoder input channels

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