test_conv.py 32.8 KB
Newer Older
yan.yan's avatar
yan.yan committed
1
2
# Copyright 2021 Yan Yan
# 
traveller59's avatar
traveller59 committed
3
4
5
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
yan.yan's avatar
yan.yan committed
6
# 
traveller59's avatar
traveller59 committed
7
#     http://www.apache.org/licenses/LICENSE-2.0
yan.yan's avatar
yan.yan committed
8
# 
traveller59's avatar
traveller59 committed
9
10
11
12
13
14
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

15
16
import time
import unittest
traveller59's avatar
traveller59 committed
17
from pathlib import Path
18
19

import numpy as np
traveller59's avatar
traveller59 committed
20
import torch
21
from torch import nn
yan.yan's avatar
v2.1  
yan.yan committed
22
from spconv.core import ConvAlgo
23

yan.yan's avatar
yan.yan committed
24
import spconv.pytorch as spconv
25
from spconv.test_utils import TestCase, generate_sparse_data, params_grid
yan.yan's avatar
yan.yan committed
26
from spconv.constants import FILTER_HWIO
27
28
# import sparseconvnet as scn

yan.yan's avatar
v2.1  
yan.yan committed
29
30
31
# we must disable tf32 to increase reference precision.
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
traveller59's avatar
traveller59 committed
32
33

class SparseConv3dTestTorch(nn.Module):
yanyan's avatar
yanyan committed
34
35
36
37
38
39
40
41
42
43
    def __init__(self,
                 num_layers,
                 ndim,
                 shape,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride,
                 padding,
                 dilation,
yan.yan's avatar
v2.1  
yan.yan committed
44
                 algo=spconv.ConvAlgo.MaskSplitImplicitGemm):
traveller59's avatar
traveller59 committed
45
        super().__init__()
yan.yan's avatar
v2.1  
yan.yan committed
46
        self.algo = algo
47
48
49
50
51
52
53
54
        layers = [
            spconv.SparseConv3d(in_channels,
                                out_channels,
                                kernel_size,
                                stride,
                                padding=padding,
                                dilation=dilation,
                                bias=False,
Yan Yan's avatar
Yan Yan committed
55
                                algo=algo)
56
        ]
traveller59's avatar
traveller59 committed
57
        for i in range(1, num_layers):
58
59
60
61
62
63
64
            layers.append(
                spconv.SparseConv3d(out_channels,
                                    out_channels,
                                    kernel_size,
                                    stride,
                                    padding=padding,
                                    dilation=dilation,
65
                                    bias=False,
Yan Yan's avatar
Yan Yan committed
66
                                    algo=algo))
67
        self.net = spconv.SparseSequential(*layers, )
traveller59's avatar
traveller59 committed
68
69
70
71
72
73
        # self.grid = torch.full([3, *shape], -1, dtype=torch.int32).cuda()
        self.grid = None
        self.shape = shape

    def forward(self, features, coors, batch_size):
        coors = coors.int()
74
75
76
77
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size,
                                    self.grid)
        return self.net(x)  # .dense()

traveller59's avatar
traveller59 committed
78
79

class SubMConv3dTestTorch(nn.Module):
yanyan's avatar
yanyan committed
80
81
82
83
84
85
86
87
88
89
90
    def __init__(self,
                 num_layers,
                 ndim,
                 shape,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride,
                 padding,
                 dilation,
                 algo=spconv.ConvAlgo.Native):
traveller59's avatar
traveller59 committed
91
        super().__init__()
92
93
94
95
96
97
98
        layers = [
            spconv.SubMConv3d(in_channels,
                              out_channels,
                              kernel_size,
                              stride,
                              padding=padding,
                              dilation=dilation,
Yan Yan's avatar
Yan Yan committed
99
100
                              bias=False,
                              algo=algo)
101
        ]
traveller59's avatar
traveller59 committed
102
        for i in range(1, num_layers):
103
104
105
106
107
108
109
            layers.append(
                spconv.SubMConv3d(out_channels,
                                  out_channels,
                                  kernel_size,
                                  stride,
                                  padding=padding,
                                  dilation=dilation,
Yan Yan's avatar
Yan Yan committed
110
                                  bias=False,
yanyan's avatar
yanyan committed
111
                                  algo=algo))
112
        self.net = spconv.SparseSequential(*layers, )
113
114
        # self.grid = torch.full([3, *shape], -1, dtype=torch.int32).cuda()
        self.grid = None
traveller59's avatar
traveller59 committed
115
116
117
        self.shape = shape

    def forward(self, features, coors, batch_size):
yanyan's avatar
yanyan committed
118
        coors = coors.int()  # .cpu()
119
120
121
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size,
                                    self.grid)
        return self.net(x)  # .dense()
traveller59's avatar
traveller59 committed
122
123
124


class Conv3dTestTorch(nn.Module):
125
126
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride, padding, dilation):
traveller59's avatar
traveller59 committed
127
        super().__init__()
128
129
130
131
132
133
134
135
136
        layers = [
            nn.Conv3d(in_channels,
                      out_channels,
                      kernel_size,
                      stride,
                      padding=padding,
                      dilation=dilation,
                      bias=False)
        ]
traveller59's avatar
traveller59 committed
137
        for i in range(1, num_layers):
138
139
140
141
142
143
144
145
146
            layers.append(
                nn.Conv3d(out_channels,
                          out_channels,
                          kernel_size,
                          stride,
                          padding=padding,
                          dilation=dilation,
                          bias=False))
        self.net = nn.Sequential(*layers, )
traveller59's avatar
traveller59 committed
147
148
149
        self.shape = shape

    def forward(self, x):
150
        return self.net(x)  # .dense()
traveller59's avatar
traveller59 committed
151
152
153


class SparseDeConv3dTestTorch(nn.Module):
154
155
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride, padding, dilation):
traveller59's avatar
traveller59 committed
156
        super().__init__()
157
158
159
160
161
162
163
164
165
        layers = [
            spconv.SparseConvTranspose3d(in_channels,
                                         out_channels,
                                         kernel_size,
                                         stride,
                                         padding=padding,
                                         dilation=dilation,
                                         bias=False)
        ]
traveller59's avatar
traveller59 committed
166
        for i in range(1, num_layers):
167
168
169
170
171
172
173
174
175
            layers.append(
                spconv.SparseConvTranspose3d(out_channels,
                                             out_channels,
                                             kernel_size,
                                             stride,
                                             padding=padding,
                                             dilation=dilation,
                                             bias=False))
        self.net = spconv.SparseSequential(*layers, )
traveller59's avatar
traveller59 committed
176
177
178
179
        self.shape = shape

    def forward(self, features, coors, batch_size):
        coors = coors.int()
180
181
182
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size)
        return self.net(x)  # .dense()

traveller59's avatar
traveller59 committed
183
184

class DeConv3dTestTorch(nn.Module):
185
186
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride, padding, dilation):
traveller59's avatar
traveller59 committed
187
        super().__init__()
188
189
190
191
192
193
194
195
196
        layers = [
            nn.ConvTranspose3d(in_channels,
                               out_channels,
                               kernel_size,
                               stride,
                               padding=padding,
                               dilation=dilation,
                               bias=False)
        ]
traveller59's avatar
traveller59 committed
197
        for i in range(1, num_layers):
198
199
200
201
202
203
204
205
206
            layers.append(
                nn.ConvTranspose3d(out_channels,
                                   out_channels,
                                   kernel_size,
                                   stride,
                                   padding=padding,
                                   dilation=dilation,
                                   bias=False))
        self.net = nn.Sequential(*layers, )
traveller59's avatar
traveller59 committed
207
208
209
        self.shape = shape

    def forward(self, x):
210
        return self.net(x)  # .dense()
traveller59's avatar
traveller59 committed
211
212
213


class SparseMaxPoolTestTorch(nn.Module):
214
215
    def __init__(self, num_layers, ndim, shape, kernel_size, stride, padding,
                 dilation):
traveller59's avatar
traveller59 committed
216
        super().__init__()
217
218
219
        layers = [
            spconv.SparseMaxPool3d(kernel_size, stride, padding, dilation)
        ]
traveller59's avatar
traveller59 committed
220
        for i in range(1, num_layers):
221
222
223
            layers.append(
                spconv.SparseMaxPool3d(kernel_size, stride, padding, dilation))
        self.net = spconv.SparseSequential(*layers, )
traveller59's avatar
traveller59 committed
224
225
226
227
        self.shape = shape

    def forward(self, features, coors, batch_size):
        coors = coors.int()
228
229
230
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size)
        return self.net(x)  # .dense()

traveller59's avatar
traveller59 committed
231
232

class MaxPool3dTestTorch(nn.Module):
233
234
    def __init__(self, num_layers, ndim, shape, kernel_size, stride, padding,
                 dilation):
traveller59's avatar
traveller59 committed
235
        super().__init__()
236
        layers = [nn.MaxPool3d(kernel_size, stride, padding, dilation)]
traveller59's avatar
traveller59 committed
237
        for i in range(1, num_layers):
238
239
            layers.append(nn.MaxPool3d(kernel_size, stride, padding, dilation))
        self.net = nn.Sequential(*layers, )
traveller59's avatar
traveller59 committed
240
241
242
        self.shape = shape

    def forward(self, x):
243
        return self.net(x)  # .dense()
traveller59's avatar
traveller59 committed
244
245
246


class SubmanifoldConvTestTorch(nn.Module):
247
248
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride):
traveller59's avatar
traveller59 committed
249
        super().__init__()
250
251
252
253
254
255
256
        layers = [
            spconv.SubMConv3d(in_channels,
                              out_channels,
                              kernel_size,
                              bias=False,
                              indice_key="subm0")
        ]
traveller59's avatar
traveller59 committed
257
        for i in range(1, num_layers):
258
259
260
261
262
263
            layers.append(
                spconv.SubMConv3d(out_channels,
                                  out_channels,
                                  kernel_size,
                                  bias=False))
        self.net = nn.Sequential(*layers, )
traveller59's avatar
traveller59 committed
264
265
266
267
        self.shape = shape

    def forward(self, features, coors, batch_size):
        coors = coors.int()
268
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size)
traveller59's avatar
traveller59 committed
269
270
        return self.net(x)

271

traveller59's avatar
traveller59 committed
272
class SCNCoupleDeConvTest(nn.Module):
273
274
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride):
traveller59's avatar
traveller59 committed
275
276
277
        super().__init__()
        self.scn_input = scn.InputLayer(ndim, shape, mode=0)
        self.net = nn.Sequential(
278
279
280
281
282
283
284
285
286
287
288
289
            scn.Convolution(ndim,
                            in_channels,
                            out_channels,
                            kernel_size,
                            stride,
                            bias=False),
            scn.Deconvolution(ndim,
                              out_channels,
                              in_channels,
                              kernel_size,
                              stride,
                              bias=False),
traveller59's avatar
traveller59 committed
290
291
292
293
294
295
296
297
            scn.SparseToDense(ndim, in_channels),
        )

    def forward(self, features, coors, batch_size):
        coors = coors.long().cpu()
        x = self.scn_input((coors, features))
        return self.net(x)

298

traveller59's avatar
traveller59 committed
299
class SparseCoupleDeConvTest(nn.Module):
300
301
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride):
traveller59's avatar
traveller59 committed
302
303
        super().__init__()
        self.net = spconv.SparseSequential(
304
305
306
307
308
309
310
311
312
313
314
            spconv.SparseConv3d(in_channels,
                                out_channels,
                                kernel_size,
                                stride,
                                indice_key="cp0",
                                bias=False),
            spconv.SparseInverseConv3d(out_channels,
                                       in_channels,
                                       kernel_size,
                                       indice_key="cp0",
                                       bias=False),
traveller59's avatar
traveller59 committed
315
316
317
318
319
320
        )
        self.todense = spconv.ToDense()
        self.shape = shape

    def forward(self, features, coors, batch_size):
        coors = coors.int()
321
322
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size)
        return self.todense(self.net(x))  # .dense()
traveller59's avatar
traveller59 committed
323
324
325
326
327


def gather_nd(params, indices):
    # this function has a limit that MAX_ADVINDEX_CALC_DIMS=5
    ndim = indices.shape[-1]
328
329
    output_shape = list(indices.shape[:-1]) + list(
        params.shape[indices.shape[-1]:])
traveller59's avatar
traveller59 committed
330
331
332
333
334
    flatted_indices = indices.view(-1, ndim)
    slices = [flatted_indices[:, i] for i in range(ndim)]
    slices += [Ellipsis]
    return params[slices].view(*output_shape)

335

traveller59's avatar
traveller59 committed
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
def scatter_nd(indices, updates, shape):
    """pytorch edition of tensorflow scatter_nd.
    this function don't contain except handle code. so use this carefully
    when indice repeats, don't support repeat add which is supported
    in tensorflow.
    """
    ret = torch.zeros(*shape, dtype=updates.dtype, device=updates.device)
    ndim = indices.shape[-1]
    output_shape = list(indices.shape[:-1]) + shape[indices.shape[-1]:]
    flatted_indices = indices.view(-1, ndim)
    slices = [flatted_indices[:, i] for i in range(ndim)]
    slices += [Ellipsis]
    ret[slices] = updates.view(*output_shape)
    return ret


class TestSpConv(TestCase):
    def testSpConv3d(self):
        np.random.seed(484)
yan.yan's avatar
v2.1  
yan.yan committed
355
        torch.manual_seed(48848)
yanyan's avatar
yanyan committed
356
        devices = ["cuda:0"]
traveller59's avatar
traveller59 committed
357
358
359
        shapes = [[19, 18, 17]]
        batchsizes = [1, 2]

yanyan's avatar
yanyan committed
360
        in_channels = [32]
traveller59's avatar
traveller59 committed
361
362
363
364
365
        out_channels = [32, 48, 64]
        ksizes = [2, 3]
        strides = [1, 2, 3]
        paddings = [0, 1, 2]
        dilations = [1, 2, 3]
yan.yan's avatar
v2.1  
yan.yan committed
366
367
        algos = [ConvAlgo.Native, ConvAlgo.MaskImplicitGemm, ConvAlgo.MaskSplitImplicitGemm]
        algos = [ConvAlgo.MaskSplitImplicitGemm]
traveller59's avatar
traveller59 committed
368

yan.yan's avatar
v2.1  
yan.yan committed
369
        for dev, shape, bs, IC, OC, k, s, p, d, al in params_grid(
370
                devices, shapes, batchsizes, in_channels, out_channels, ksizes,
yan.yan's avatar
v2.1  
yan.yan committed
371
                strides, paddings, dilations, algos):
traveller59's avatar
traveller59 committed
372
            if all([s > 1, d > 1]):
373
                continue  # don't support this.
yan.yan's avatar
v2.1  
yan.yan committed
374
            print(k, s, p, d)
traveller59's avatar
traveller59 committed
375
376
            device = torch.device(dev)
            num_points = [1000] * bs
yan.yan's avatar
v2.1  
yan.yan committed
377
378
379
380
381
382
            dtype = torch.float32
            net = SparseConv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
                                        d, algo=al).to(device).to(dtype)
            net_ref = Conv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
                                      d).to(device).to(dtype)

traveller59's avatar
traveller59 committed
383
384
            sparse_dict = generate_sparse_data(shape, num_points, IC)

385
386
387
388
            features = np.ascontiguousarray(sparse_dict["features"]).astype(
                np.float32)
            indices = np.ascontiguousarray(
                sparse_dict["indices"][:, [3, 0, 1, 2]]).astype(np.int32)
traveller59's avatar
traveller59 committed
389
390
            features_dense = sparse_dict["features_dense"].astype(np.float32)
            indices_t = torch.from_numpy(indices).int().to(device)
yan.yan's avatar
bug fix  
yan.yan committed
391
            features_t = torch.from_numpy(features).to(device).to(dtype)
traveller59's avatar
traveller59 committed
392
            features_t.requires_grad = True
yan.yan's avatar
bug fix  
yan.yan committed
393
            features_dense_t = torch.from_numpy(features_dense).to(device).to(dtype)
traveller59's avatar
traveller59 committed
394
            features_dense_t.requires_grad = True
yan.yan's avatar
v2.1  
yan.yan committed
395
396
397
398
399
400
401
402
403
404
405
406
407
408
            if net.algo == ConvAlgo.Native:
                if FILTER_HWIO:
                    filters = np.random.uniform(-1, 1, size=[k, k, k, IC,
                                                            OC]).astype(np.float32)
                else:
                    filters = np.random.uniform(-1, 1, size=[k, k, k, OC,
                                                            IC]).astype(np.float32)
                filters_t = torch.from_numpy(filters).to(device).to(dtype)
                if FILTER_HWIO:
                    net_ref.net[0].weight.data[:] = filters_t.permute(4, 3, 0, 1,
                                                                    2).contiguous()
                else:
                    net_ref.net[0].weight.data[:] = filters_t.permute(3, 4, 0, 1,
                                                                    2).contiguous()
yan.yan's avatar
yan.yan committed
409
            else:
yan.yan's avatar
v2.1  
yan.yan committed
410
411
412
413
414
                filters = np.random.uniform(-1, 1, size=[OC, k, k, k, IC]).astype(np.float32)
                filters_t = torch.from_numpy(filters).to(device).to(dtype)
                net_ref.net[0].weight.data[:] = filters_t.permute(0, 4, 1, 2,
                                                                3).contiguous()

traveller59's avatar
traveller59 committed
415
416
            net.net[0].weight.data[:] = filters_t
            out_ref = net_ref(features_dense_t)
417
            out = net(features_t, indices_t, bs).dense()
yan.yan's avatar
yan.yan committed
418
419
420
421
            out_np = out.detach().cpu().numpy()
            out_ref_np = out_ref.detach().cpu().numpy()
            self.assertAllClose(out_np, out_ref_np, atol=1e-4)

422
423
            dout = np.random.uniform(-0.2, 0.2,
                                     out_ref.shape).astype(features.dtype)
traveller59's avatar
traveller59 committed
424
425
426
            dout_t = torch.from_numpy(dout).to(device)
            out.backward(dout_t)
            out_ref.backward(dout_t)
427
428
            din_dense = features_dense_t.grad.detach().permute(0, 2, 3, 4,
                                                               1).contiguous()
traveller59's avatar
traveller59 committed
429
430
            din_sparse = gather_nd(din_dense, indices_t.long())
            din = features_t.grad.detach()
yan.yan's avatar
yan.yan committed
431

traveller59's avatar
traveller59 committed
432
433
434
435
436
            din_np = din.cpu().numpy()
            din_sparse_np = din_sparse.cpu().numpy()
            for layer, layer_ref in zip(net.net, net_ref.net):
                dw = layer.weight.grad.detach().cpu().numpy()
                dw_ref = layer_ref.weight.grad.detach().cpu().numpy()
yan.yan's avatar
v2.1  
yan.yan committed
437
438
439
440
441
                if net.algo == ConvAlgo.Native:
                    if FILTER_HWIO:
                        dw = dw.transpose(4, 3, 0, 1, 2)
                    else:
                        dw = dw.transpose(3, 4, 0, 1, 2)
yan.yan's avatar
yan.yan committed
442
                else:
yan.yan's avatar
v2.1  
yan.yan committed
443
444
                    # OHWI -> OIHW
                    dw = dw.transpose(0, 4, 1, 2, 3)
yan.yan's avatar
yan.yan committed
445

traveller59's avatar
traveller59 committed
446
                self.assertAllClose(dw, dw_ref, atol=1e-4)
yan.yan's avatar
yan.yan committed
447
            self.assertAllClose(din_np, din_sparse_np, atol=1e-4)
traveller59's avatar
traveller59 committed
448

449

traveller59's avatar
traveller59 committed
450
451
    def testSpDeConv3d(self):
        np.random.seed(484)
Yan Yan's avatar
Yan Yan committed
452
        devices = ["cuda:0"]
traveller59's avatar
traveller59 committed
453
454
455
456
457
458
459
460
461
        shapes = [[19, 18, 17]]
        batchsizes = [1, 2]

        in_channels = [64]
        out_channels = [32, 48, 64]
        ksizes = [2, 3]
        strides = [2, 3]
        paddings = [0, 1, 2]
        dilations = [1, 2, 3]
yan.yan's avatar
bug fix  
yan.yan committed
462
463
464
465
466
        ksizes = [3]

        strides = [1]
        paddings = [0]
        dilations = [1]
traveller59's avatar
traveller59 committed
467
468

        for dev, shape, bs, IC, OC, k, s, p, d in params_grid(
469
470
                devices, shapes, batchsizes, in_channels, out_channels, ksizes,
                strides, paddings, dilations):
traveller59's avatar
traveller59 committed
471
            if all([s > 1, d > 1]):
472
                continue  # don't support this.
traveller59's avatar
traveller59 committed
473
474
475
476
477
            device = torch.device(dev)
            num_points = [1000] * bs

            sparse_dict = generate_sparse_data(shape, num_points, IC)

478
479
480
481
            features = np.ascontiguousarray(sparse_dict["features"]).astype(
                np.float32)
            indices = np.ascontiguousarray(
                sparse_dict["indices"][:, [3, 0, 1, 2]]).astype(np.int32)
traveller59's avatar
traveller59 committed
482
            features_dense = sparse_dict["features_dense"].astype(np.float32)
yan.yan's avatar
bug fix  
yan.yan committed
483
484
485
486
487
488
489
            if FILTER_HWIO:
                filters = np.random.uniform(0, 1, size=[k, k, k, IC,
                                                        OC]).astype(np.float32)
            else:
                filters = np.random.uniform(0, 1, size=[k, k, k, OC,
                                                        IC]).astype(np.float32)

traveller59's avatar
traveller59 committed
490
491
492
493
494
            indices_t = torch.from_numpy(indices).int().to(device)
            features_t = torch.from_numpy(features).to(device)
            features_t.requires_grad = True
            features_dense_t = torch.from_numpy(features_dense).to(device)
            features_dense_t.requires_grad = True
495
496
497
498
            net = SparseDeConv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
                                          d).to(device)
            net_ref = DeConv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
                                        d).to(device)
traveller59's avatar
traveller59 committed
499
            filters_t = torch.from_numpy(filters).to(device)
yan.yan's avatar
bug fix  
yan.yan committed
500
501
502
503
504
505
506
            print(net_ref.net[0].weight.shape)
            if FILTER_HWIO:
                net_ref.net[0].weight.data[:] = filters_t.permute(3, 4, 0, 1,
                                                                2).contiguous()
            else:
                net_ref.net[0].weight.data[:] = filters_t.permute(4, 3, 0, 1,
                                                                2).contiguous()
traveller59's avatar
traveller59 committed
507
508
            net.net[0].weight.data[:] = filters_t
            out_ref = net_ref(features_dense_t)
509
            out = net(features_t, indices_t, bs).dense()
yan.yan's avatar
bug fix  
yan.yan committed
510
511
512
513
            out_np = out.detach().cpu().numpy()
            out_ref_np = out_ref.detach().cpu().numpy()
            self.assertAllClose(out_np, out_ref_np, atol=1e-4)

514
515
            dout = np.random.uniform(-0.2, 0.2,
                                     out_ref.shape).astype(features.dtype)
traveller59's avatar
traveller59 committed
516
517
518
            dout_t = torch.from_numpy(dout).to(device)
            out.backward(dout_t)
            out_ref.backward(dout_t)
519
520
            din_dense = features_dense_t.grad.detach().permute(0, 2, 3, 4,
                                                               1).contiguous()
traveller59's avatar
traveller59 committed
521
522
523
524
525
526
527
528
            din_sparse = gather_nd(din_dense, indices_t.long())
            din = features_t.grad.detach()
            din_np = din.cpu().numpy()
            din_sparse_np = din_sparse.cpu().numpy()
            self.assertAllClose(din_np, din_sparse_np, atol=1e-4)
            for layer, layer_ref in zip(net.net, net_ref.net):
                dw = layer.weight.grad.detach().cpu().numpy()
                dw_ref = layer_ref.weight.grad.detach().cpu().numpy()
yan.yan's avatar
bug fix  
yan.yan committed
529
530
531
532
                if FILTER_HWIO:
                    dw = dw.transpose(3, 4, 0, 1, 2)
                else:
                    dw = dw.transpose(4, 3, 0, 1, 2)
traveller59's avatar
traveller59 committed
533
534
                self.assertAllClose(dw, dw_ref, atol=1e-4)

535

traveller59's avatar
traveller59 committed
536
537
538
539
540
541
542
543
544
545
546
547
548
549
    def testSpCpConv3d(self):
        np.random.seed(484)
        devices = ["cuda:0", "cpu:0"]
        shapes = [[20, 20, 20]]
        batchsizes = [1, 2]

        in_channels = [64]
        out_channels = [32, 48, 64]
        ksizes = [2]
        strides = [2]
        paddings = [0, 1, 2]
        dilations = [1, 2, 3]

        for dev, shape, bs, IC, OC, k, s in params_grid(
550
551
                devices, shapes, batchsizes, in_channels, out_channels, ksizes,
                strides):
traveller59's avatar
traveller59 committed
552
553
554
555
556
            device = torch.device(dev)
            num_points = [1000] * bs

            sparse_dict = generate_sparse_data(shape, num_points, IC)

557
558
559
560
            features = np.ascontiguousarray(sparse_dict["features"]).astype(
                np.float32)
            indices = np.ascontiguousarray(
                sparse_dict["indices"][:, [3, 0, 1, 2]]).astype(np.int32)
traveller59's avatar
traveller59 committed
561
            features_dense = sparse_dict["features_dense"].astype(np.float32)
562
563
            filters = np.random.uniform(0, 1, size=[k, k, k, IC,
                                                    OC]).astype(np.float32)
traveller59's avatar
traveller59 committed
564
            indices_t = torch.from_numpy(indices).int().to(device)
565
566
            indices_scn_t = torch.from_numpy(
                indices[:, [1, 2, 3, 0]]).int().to(device)
traveller59's avatar
traveller59 committed
567
568
569
570
571
572
573
            features_t = torch.from_numpy(features).to(device)
            features_t.requires_grad = True
            features_ref_t = torch.from_numpy(features).to(device)
            features_ref_t.requires_grad = True

            net_ref = SCNCoupleDeConvTest(1, 3, shape, IC, OC, k, s).to(device)
            net = SparseCoupleDeConvTest(1, 3, shape, IC, OC, k, s).to(device)
574
575
576
577
            net_ref.net[0].weight.data[:] = net.net[0].weight.data[:].view(
                *net_ref.net[0].weight.shape)
            net_ref.net[1].weight.data[:] = net.net[1].weight.data[:].view(
                *net_ref.net[1].weight.shape)
traveller59's avatar
traveller59 committed
578
579
            out_ref = net_ref(features_ref_t, indices_scn_t, bs)
            out = net(features_t, indices_t, bs)
580
581
            dout = np.random.uniform(-0.2, 0.2,
                                     out_ref.shape).astype(features.dtype)
traveller59's avatar
traveller59 committed
582
583
584
585
586
587
588
589
590
591
            dout_t = torch.from_numpy(dout).to(device)
            out.backward(dout_t)
            out_ref.backward(dout_t)
            din = features_t.grad.detach()
            din_ref = features_ref_t.grad.detach()
            din_np = din.cpu().numpy()
            din_ref_np = din_ref.cpu().numpy()
            self.assertAllClose(din_ref_np, din_np, atol=1e-4)
            for layer, layer_ref in zip(net.net, net_ref.net):
                dw = layer.weight.grad.detach().cpu().numpy()
592
593
                dw_ref = layer_ref.weight.grad.detach().cpu().view(
                    *dw.shape).numpy()
traveller59's avatar
traveller59 committed
594
595
596
597
598
599
600
                self.assertAllClose(dw, dw_ref, atol=1e-4)

            out_np = out.detach().cpu().numpy()
            out_ref_np = out_ref.detach().cpu().numpy()
            self.assertAllClose(out_np, out_ref_np, atol=1e-4)

    def testSpMaxPool3d(self):
601
        np.random.seed(485)
602
        devices = ["cuda:0"]
traveller59's avatar
traveller59 committed
603
604
605
        shapes = [[19, 18, 17]]
        batchsizes = [1, 2]

yan.yan's avatar
yan.yan committed
606
607
        in_channels = [64]
        out_channels = [64]
traveller59's avatar
traveller59 committed
608
609
610
611
        ksizes = [2, 3]
        strides = [1, 2, 3]
        paddings = [0, 1]
        dilations = [1, 2, 3]
yan.yan's avatar
v2.1  
yan.yan committed
612
613
614
615
        # ksizes = [2]
        # strides = [2]
        # paddings = [0]
        # dilations = [1]
traveller59's avatar
traveller59 committed
616
617

        for dev, shape, bs, IC, OC, k, s, p, d in params_grid(
618
619
                devices, shapes, batchsizes, in_channels, out_channels, ksizes,
                strides, paddings, dilations):
traveller59's avatar
traveller59 committed
620
            if all([s > 1, d > 1]):
621
                continue  # don't support this.
traveller59's avatar
traveller59 committed
622
623
            device = torch.device(dev)
            num_points = [1000] * bs
yan.yan's avatar
yan.yan committed
624

traveller59's avatar
traveller59 committed
625
            # when data contains negative, sparse maxpool is not equal to dense maxpool.
626
627
628
629
630
631
632
633
634
            sparse_dict = generate_sparse_data(shape,
                                               num_points,
                                               IC,
                                               data_range=[0.1, 1])

            features = np.ascontiguousarray(sparse_dict["features"]).astype(
                np.float32)
            indices = np.ascontiguousarray(
                sparse_dict["indices"][:, [3, 0, 1, 2]]).astype(np.int32)
traveller59's avatar
traveller59 committed
635
            features_dense = sparse_dict["features_dense"].astype(np.float32)
yan.yan's avatar
yan.yan committed
636
637
            filters = np.random.uniform(0, 1, size=[k, k, k, OC,
                                                    IC]).astype(np.float32)
traveller59's avatar
traveller59 committed
638
639
640
641
642
643
644
645
646
647
            indices_t = torch.from_numpy(indices).int().to(device)
            features_t = torch.from_numpy(features).to(device)
            features_t.requires_grad = True
            features_dense_t = torch.from_numpy(features_dense).to(device)
            features_dense_t.requires_grad = True
            net = SparseMaxPoolTestTorch(1, 3, shape, k, s, p, d).to(device)
            net_ref = MaxPool3dTestTorch(1, 3, shape, k, s, p, d).to(device)

            out_ref = net_ref(features_dense_t)
            out = net(features_t, indices_t, bs)
yan.yan's avatar
yan.yan committed
648

traveller59's avatar
traveller59 committed
649
650
            outids = out.indices
            outfeatures = out.features
651
            outids_dev = outids.float()
652
            out_dense = out.dense(channels_first=False)
traveller59's avatar
traveller59 committed
653
            out = out_dense.permute(0, 4, 1, 2, 3).contiguous()
yan.yan's avatar
yan.yan committed
654
655
656
            out_np = out.detach().cpu().numpy()
            out_ref_np = out_ref.detach().cpu().numpy()
            self.assertAllClose(out_np, out_ref_np, atol=1e-4)
657
658
659

            dout_sparse = np.random.uniform(
                -0.2, 0.2, outfeatures.shape).astype(features.dtype)
traveller59's avatar
traveller59 committed
660
            dout_sparse_t = torch.from_numpy(dout_sparse).to(device)
661
662
            dout_t = scatter_nd(outids.long(), dout_sparse_t,
                                list(out_dense.shape))
traveller59's avatar
traveller59 committed
663
664
665
            dout_t = dout_t.permute(0, 4, 1, 2, 3).contiguous()
            out.backward(dout_t)
            out_ref.backward(dout_t)
666
667
            din_dense = features_dense_t.grad.detach().permute(0, 2, 3, 4,
                                                               1).contiguous()
traveller59's avatar
traveller59 committed
668
669
            din_sparse = gather_nd(din_dense, indices_t.long())
            din = features_t.grad.detach()
670

671
672
673
            din_np = din.cpu().numpy()
            din_sparse_np = din_sparse.cpu().numpy()
            self.assertAllClose(din_np, din_sparse_np, atol=1e-4)
674

traveller59's avatar
traveller59 committed
675

yanyan's avatar
yanyan committed
676
def main(algo=spconv.ConvAlgo.Native, dtype=torch.float32):
traveller59's avatar
traveller59 committed
677
678
    # function for develop.
    np.random.seed(484)
traveller59's avatar
traveller59 committed
679
    # devices = ["cuda:0"]
traveller59's avatar
traveller59 committed
680
    devices = ["cuda:0"]
yanyan's avatar
yanyan committed
681
    shapes = [[400, 400, 15]]
yanyan's avatar
yanyan committed
682
    batchsizes = [2]
traveller59's avatar
traveller59 committed
683

yanyan's avatar
yanyan committed
684
685
    in_channels = [19]
    out_channels = [17]
Yan Yan's avatar
Yan Yan committed
686
    ksizes = [(3, 3, 3)]
traveller59's avatar
traveller59 committed
687
688
689
690
691
    strides = [1]
    paddings = [0]
    dilations = [1]

    for dev, shape, bs, IC, OC, k, s, p, d in params_grid(
692
693
            devices, shapes, batchsizes, in_channels, out_channels, ksizes,
            strides, paddings, dilations):
traveller59's avatar
traveller59 committed
694
695
696
        if all([s > 1, d > 1]):
            continue
        device = torch.device(dev)
yanyan's avatar
yanyan committed
697
        num_points = [30000] * bs
traveller59's avatar
traveller59 committed
698
699
700

        sparse_dict = generate_sparse_data(shape, num_points, IC)

701
702
703
704
        features = np.ascontiguousarray(sparse_dict["features"]).astype(
            np.float32)
        indices = np.ascontiguousarray(
            sparse_dict["indices"][:, [3, 0, 1, 2]]).astype(np.int32)
traveller59's avatar
traveller59 committed
705
706
        features_dense = sparse_dict["features_dense"].astype(np.float32)
        indices_t = torch.from_numpy(indices)
707
708
        filters = np.random.uniform(0, 1, size=[k[0], 1, 1, IC,
                                                OC]).astype(np.float32)
yanyan's avatar
yanyan committed
709
710
        indices_t = torch.from_numpy(indices).int().to(device).to(dtype)
        features_t = torch.from_numpy(features).to(device).to(dtype)
711

yanyan's avatar
yanyan committed
712
713
        features_dense_t = torch.from_numpy(features_dense).to(device).to(
            dtype)
yanyan's avatar
yanyan committed
714
        net = SparseConv3dTestTorch(1, 3, shape, IC, OC, k, s, p, d,
yanyan's avatar
yanyan committed
715
                                    algo=algo).to(device).to(dtype)
716
        net_ref = Conv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
yanyan's avatar
yanyan committed
717
718
                                  d).to(device).to(dtype)
        filters_t = torch.from_numpy(filters).to(device).to(dtype)
719
720
        net_ref.net[0].weight[:] = filters_t.permute(4, 3, 0, 1,
                                                     2).contiguous()
traveller59's avatar
traveller59 committed
721
722
723
        net.net[0].weight[:] = filters_t
        out_ref = net_ref(features_dense_t)
        times = []
tusimple's avatar
tusimple committed
724
        for i in range(10):
traveller59's avatar
traveller59 committed
725
726
727
728
            t = time.time()
            out = net(features_t, indices_t, bs)
            torch.cuda.synchronize()
            times.append(time.time() - t)
729
        # print((net.grid == -1).float().sum(), net.grid.numel())
730
        # print("spconv time", time.time() - t)
traveller59's avatar
traveller59 committed
731
        print("spconv time", np.mean(times[2:]))
traveller59's avatar
traveller59 committed
732
733
734
        out = net(features_t, indices_t, bs)
        # print(out.indices)
        out = out.dense()
Yan Yan's avatar
Yan Yan committed
735
736
        out_numpy = out.detach().cpu().numpy()

737
738
739
        print(
            np.linalg.norm(out.detach().cpu().numpy() -
                           out_ref.detach().cpu().numpy()))
yanyan's avatar
yanyan committed
740
741
        print(out_numpy.min(), out_numpy.max(), out_numpy.mean(),
              out_numpy.sum())
traveller59's avatar
traveller59 committed
742
743


yanyan's avatar
yanyan committed
744
def main_subm(algo, dtype=torch.float32):
745
746
    # function for develop.
    np.random.seed(484)
Yan Yan's avatar
Yan Yan committed
747
    torch.manual_seed(50051)
748
749
    # devices = ["cuda:0"]
    devices = ["cuda:0"]
yanyan's avatar
yanyan committed
750
    shapes = [[400, 400, 15]]
751
752
753
754
755
756
757
758
759
760
761
762
763
764
    batchsizes = [2]

    in_channels = [32]
    out_channels = [64]
    ksizes = [(3, 3, 3)]
    strides = [1]
    paddings = [1]
    dilations = [1]
    for dev, shape, bs, IC, OC, k, s, p, d in params_grid(
            devices, shapes, batchsizes, in_channels, out_channels, ksizes,
            strides, paddings, dilations):
        if all([s > 1, d > 1]):
            continue
        device = torch.device(dev)
yanyan's avatar
yanyan committed
765
        num_points = [120000] * bs
766
767
768
769
770
771
772
773
774
775
776

        sparse_dict = generate_sparse_data(shape, num_points, IC)

        features = np.ascontiguousarray(sparse_dict["features"]).astype(
            np.float32)
        indices = np.ascontiguousarray(
            sparse_dict["indices"][:, [3, 0, 1, 2]]).astype(np.int32)
        features_dense = sparse_dict["features_dense"].astype(np.float32)
        indices_t = torch.from_numpy(indices)
        filters = np.random.uniform(0, 1, size=[k[0], 1, 1, IC,
                                                OC]).astype(np.float32)
yanyan's avatar
yanyan committed
777
778
        indices_t = torch.from_numpy(indices).int().to(device).to(dtype)
        features_t = torch.from_numpy(features).to(device).to(dtype)
779

yanyan's avatar
yanyan committed
780
781
        features_dense_t = torch.from_numpy(features_dense).to(device).to(
            dtype)
yanyan's avatar
yanyan committed
782
        net = SubMConv3dTestTorch(1, 3, shape, IC, OC, k, s, p, d,
yanyan's avatar
yanyan committed
783
                                  algo=algo).to(device).to(dtype)
784
        net_ref = Conv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
yanyan's avatar
yanyan committed
785
786
                                  d).to(device).to(dtype)
        filters_t = torch.from_numpy(filters).to(device).to(dtype)
787
788
789
790
791
        net_ref.net[0].weight[:] = filters_t.permute(4, 3, 0, 1,
                                                     2).contiguous()
        net.net[0].weight[:] = filters_t
        out_ref = net_ref(features_dense_t)
        times = []
yanyan's avatar
yanyan committed
792
        for i in range(20):
793
794
795
796
797
798
799
800
801
802
            t = time.time()
            out = net(features_t, indices_t, bs)
            torch.cuda.synchronize()
            times.append(time.time() - t)
        # print((net.grid == -1).float().sum(), net.grid.numel())
        # print("spconv time", time.time() - t)
        print("spconv time", np.mean(times[10:]))
        out = net(features_t, indices_t, bs)
        # print(out.indices)
        out = out.dense()
Yan Yan's avatar
Yan Yan committed
803
        out_numpy = out.detach().cpu().numpy()
yanyan's avatar
yanyan committed
804
805
806
        # print(
        #     np.linalg.norm(out.detach().cpu().numpy() -
        #                    out_ref.detach().cpu().numpy()))
yanyan's avatar
yanyan committed
807
808
        print(out_numpy.min(), out_numpy.max(), out_numpy.mean(),
              out_numpy.sum())
Yan Yan's avatar
Yan Yan committed
809
    return out_numpy
810

yanyan's avatar
yanyan committed
811

traveller59's avatar
traveller59 committed
812
if __name__ == '__main__':
yanyan's avatar
yanyan committed
813
814
    # main_subm(algo=spconv.ConvAlgo.SparseConvNet, dtype=torch.float32)
    # main(algo=spconv.ConvAlgo.SparseConvNet, dtype=torch.float32)
Yan Yan's avatar
Yan Yan committed
815
    # TestCase().assertAllClose(out_my, out_ref)
traveller59's avatar
traveller59 committed
816
    # unittest.main()
yan.yan's avatar
v2.1  
yan.yan committed
817
    TestSpConv().testSpMaxPool3d()