test_conv.py 31.9 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
22
from torch import nn

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

traveller59's avatar
traveller59 committed
28
29

class SparseConv3dTestTorch(nn.Module):
yanyan's avatar
yanyan committed
30
31
32
33
34
35
36
37
38
39
    def __init__(self,
                 num_layers,
                 ndim,
                 shape,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride,
                 padding,
                 dilation,
yan.yan's avatar
yan.yan committed
40
                 algo=spconv.ConvAlgo.Native):
traveller59's avatar
traveller59 committed
41
        super().__init__()
42
43
44
45
46
47
48
49
        layers = [
            spconv.SparseConv3d(in_channels,
                                out_channels,
                                kernel_size,
                                stride,
                                padding=padding,
                                dilation=dilation,
                                bias=False,
Yan Yan's avatar
Yan Yan committed
50
                                algo=algo)
51
        ]
traveller59's avatar
traveller59 committed
52
        for i in range(1, num_layers):
53
54
55
56
57
58
59
            layers.append(
                spconv.SparseConv3d(out_channels,
                                    out_channels,
                                    kernel_size,
                                    stride,
                                    padding=padding,
                                    dilation=dilation,
60
                                    bias=False,
Yan Yan's avatar
Yan Yan committed
61
                                    algo=algo))
62
        self.net = spconv.SparseSequential(*layers, )
traveller59's avatar
traveller59 committed
63
64
65
66
67
68
        # 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()
69
70
71
72
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size,
                                    self.grid)
        return self.net(x)  # .dense()

traveller59's avatar
traveller59 committed
73
74

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

    def forward(self, features, coors, batch_size):
yanyan's avatar
yanyan committed
113
        coors = coors.int()  # .cpu()
114
115
116
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size,
                                    self.grid)
        return self.net(x)  # .dense()
traveller59's avatar
traveller59 committed
117
118
119


class Conv3dTestTorch(nn.Module):
120
121
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride, padding, dilation):
traveller59's avatar
traveller59 committed
122
        super().__init__()
123
124
125
126
127
128
129
130
131
        layers = [
            nn.Conv3d(in_channels,
                      out_channels,
                      kernel_size,
                      stride,
                      padding=padding,
                      dilation=dilation,
                      bias=False)
        ]
traveller59's avatar
traveller59 committed
132
        for i in range(1, num_layers):
133
134
135
136
137
138
139
140
141
            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
142
143
144
        self.shape = shape

    def forward(self, x):
145
        return self.net(x)  # .dense()
traveller59's avatar
traveller59 committed
146
147
148


class SparseDeConv3dTestTorch(nn.Module):
149
150
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride, padding, dilation):
traveller59's avatar
traveller59 committed
151
        super().__init__()
152
153
154
155
156
157
158
159
160
        layers = [
            spconv.SparseConvTranspose3d(in_channels,
                                         out_channels,
                                         kernel_size,
                                         stride,
                                         padding=padding,
                                         dilation=dilation,
                                         bias=False)
        ]
traveller59's avatar
traveller59 committed
161
        for i in range(1, num_layers):
162
163
164
165
166
167
168
169
170
            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
171
172
173
174
        self.shape = shape

    def forward(self, features, coors, batch_size):
        coors = coors.int()
175
176
177
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size)
        return self.net(x)  # .dense()

traveller59's avatar
traveller59 committed
178
179

class DeConv3dTestTorch(nn.Module):
180
181
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride, padding, dilation):
traveller59's avatar
traveller59 committed
182
        super().__init__()
183
184
185
186
187
188
189
190
191
        layers = [
            nn.ConvTranspose3d(in_channels,
                               out_channels,
                               kernel_size,
                               stride,
                               padding=padding,
                               dilation=dilation,
                               bias=False)
        ]
traveller59's avatar
traveller59 committed
192
        for i in range(1, num_layers):
193
194
195
196
197
198
199
200
201
            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
202
203
204
        self.shape = shape

    def forward(self, x):
205
        return self.net(x)  # .dense()
traveller59's avatar
traveller59 committed
206
207
208


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

    def forward(self, features, coors, batch_size):
        coors = coors.int()
223
224
225
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size)
        return self.net(x)  # .dense()

traveller59's avatar
traveller59 committed
226
227

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

    def forward(self, x):
238
        return self.net(x)  # .dense()
traveller59's avatar
traveller59 committed
239
240
241


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

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

266

traveller59's avatar
traveller59 committed
267
class SCNCoupleDeConvTest(nn.Module):
268
269
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride):
traveller59's avatar
traveller59 committed
270
271
272
        super().__init__()
        self.scn_input = scn.InputLayer(ndim, shape, mode=0)
        self.net = nn.Sequential(
273
274
275
276
277
278
279
280
281
282
283
284
            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
285
286
287
288
289
290
291
292
            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)

293

traveller59's avatar
traveller59 committed
294
class SparseCoupleDeConvTest(nn.Module):
295
296
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride):
traveller59's avatar
traveller59 committed
297
298
        super().__init__()
        self.net = spconv.SparseSequential(
299
300
301
302
303
304
305
306
307
308
309
            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
310
311
312
313
314
315
        )
        self.todense = spconv.ToDense()
        self.shape = shape

    def forward(self, features, coors, batch_size):
        coors = coors.int()
316
317
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size)
        return self.todense(self.net(x))  # .dense()
traveller59's avatar
traveller59 committed
318
319
320
321
322


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

330

traveller59's avatar
traveller59 committed
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
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)
yanyan's avatar
yanyan committed
350
        devices = ["cuda:0"]
traveller59's avatar
traveller59 committed
351
352
353
        shapes = [[19, 18, 17]]
        batchsizes = [1, 2]

yanyan's avatar
yanyan committed
354
        in_channels = [32]
traveller59's avatar
traveller59 committed
355
356
357
358
359
        out_channels = [32, 48, 64]
        ksizes = [2, 3]
        strides = [1, 2, 3]
        paddings = [0, 1, 2]
        dilations = [1, 2, 3]
yan.yan's avatar
yan.yan committed
360
361
362
        # strides = [1]
        # paddings = [0]
        # dilations = [1]
traveller59's avatar
traveller59 committed
363
364

        for dev, shape, bs, IC, OC, k, s, p, d in params_grid(
365
366
                devices, shapes, batchsizes, in_channels, out_channels, ksizes,
                strides, paddings, dilations):
traveller59's avatar
traveller59 committed
367
            if all([s > 1, d > 1]):
368
                continue  # don't support this.
traveller59's avatar
traveller59 committed
369
370
371
372
            device = torch.device(dev)
            num_points = [1000] * bs
            sparse_dict = generate_sparse_data(shape, num_points, IC)

373
374
375
376
            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
377
            features_dense = sparse_dict["features_dense"].astype(np.float32)
yan.yan's avatar
yan.yan committed
378
379
380
381
382
383
            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)
yan.yan's avatar
bug fix  
yan.yan committed
384
            dtype = torch.float16
traveller59's avatar
traveller59 committed
385
            indices_t = torch.from_numpy(indices).int().to(device)
yan.yan's avatar
bug fix  
yan.yan committed
386
            features_t = torch.from_numpy(features).to(device).to(dtype)
traveller59's avatar
traveller59 committed
387
            features_t.requires_grad = True
yan.yan's avatar
bug fix  
yan.yan committed
388
            features_dense_t = torch.from_numpy(features_dense).to(device).to(dtype)
traveller59's avatar
traveller59 committed
389
            features_dense_t.requires_grad = True
390
            net = SparseConv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
yan.yan's avatar
bug fix  
yan.yan committed
391
                                        d).to(device).to(dtype)
392
            net_ref = Conv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
yan.yan's avatar
bug fix  
yan.yan committed
393
394
                                      d).to(device).to(dtype)
            filters_t = torch.from_numpy(filters).to(device).to(dtype)
yan.yan's avatar
yan.yan committed
395
396
397
398
399
400
            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()
traveller59's avatar
traveller59 committed
401
402
            net.net[0].weight.data[:] = filters_t
            out_ref = net_ref(features_dense_t)
403
            out = net(features_t, indices_t, bs).dense()
yan.yan's avatar
yan.yan committed
404
405
406
407
            out_np = out.detach().cpu().numpy()
            out_ref_np = out_ref.detach().cpu().numpy()
            self.assertAllClose(out_np, out_ref_np, atol=1e-4)

408
409
            dout = np.random.uniform(-0.2, 0.2,
                                     out_ref.shape).astype(features.dtype)
traveller59's avatar
traveller59 committed
410
411
412
            dout_t = torch.from_numpy(dout).to(device)
            out.backward(dout_t)
            out_ref.backward(dout_t)
413
414
            din_dense = features_dense_t.grad.detach().permute(0, 2, 3, 4,
                                                               1).contiguous()
traveller59's avatar
traveller59 committed
415
416
            din_sparse = gather_nd(din_dense, indices_t.long())
            din = features_t.grad.detach()
yan.yan's avatar
yan.yan committed
417

traveller59's avatar
traveller59 committed
418
419
420
421
422
            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
yan.yan committed
423
424
425
426
427
428
                if FILTER_HWIO:

                    dw = dw.transpose(4, 3, 0, 1, 2)
                else:
                    dw = dw.transpose(3, 4, 0, 1, 2)

traveller59's avatar
traveller59 committed
429
                self.assertAllClose(dw, dw_ref, atol=1e-4)
yan.yan's avatar
yan.yan committed
430
            self.assertAllClose(din_np, din_sparse_np, atol=1e-4)
traveller59's avatar
traveller59 committed
431

432

traveller59's avatar
traveller59 committed
433
434
    def testSpDeConv3d(self):
        np.random.seed(484)
Yan Yan's avatar
Yan Yan committed
435
        devices = ["cuda:0"]
traveller59's avatar
traveller59 committed
436
437
438
439
440
441
442
443
444
        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
445
446
447
448
449
        ksizes = [3]

        strides = [1]
        paddings = [0]
        dilations = [1]
traveller59's avatar
traveller59 committed
450
451

        for dev, shape, bs, IC, OC, k, s, p, d in params_grid(
452
453
                devices, shapes, batchsizes, in_channels, out_channels, ksizes,
                strides, paddings, dilations):
traveller59's avatar
traveller59 committed
454
            if all([s > 1, d > 1]):
455
                continue  # don't support this.
traveller59's avatar
traveller59 committed
456
457
458
459
460
            device = torch.device(dev)
            num_points = [1000] * bs

            sparse_dict = generate_sparse_data(shape, num_points, IC)

461
462
463
464
            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
465
            features_dense = sparse_dict["features_dense"].astype(np.float32)
yan.yan's avatar
bug fix  
yan.yan committed
466
467
468
469
470
471
472
            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
473
474
475
476
477
            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
478
479
480
481
            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
482
            filters_t = torch.from_numpy(filters).to(device)
yan.yan's avatar
bug fix  
yan.yan committed
483
484
485
486
487
488
489
            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
490
491
            net.net[0].weight.data[:] = filters_t
            out_ref = net_ref(features_dense_t)
492
            out = net(features_t, indices_t, bs).dense()
yan.yan's avatar
bug fix  
yan.yan committed
493
494
495
496
            out_np = out.detach().cpu().numpy()
            out_ref_np = out_ref.detach().cpu().numpy()
            self.assertAllClose(out_np, out_ref_np, atol=1e-4)

497
498
            dout = np.random.uniform(-0.2, 0.2,
                                     out_ref.shape).astype(features.dtype)
traveller59's avatar
traveller59 committed
499
500
501
            dout_t = torch.from_numpy(dout).to(device)
            out.backward(dout_t)
            out_ref.backward(dout_t)
502
503
            din_dense = features_dense_t.grad.detach().permute(0, 2, 3, 4,
                                                               1).contiguous()
traveller59's avatar
traveller59 committed
504
505
506
507
508
509
510
511
            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
512
513
514
515
                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
516
517
                self.assertAllClose(dw, dw_ref, atol=1e-4)

518

traveller59's avatar
traveller59 committed
519
520
521
522
523
524
525
526
527
528
529
530
531
532
    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(
533
534
                devices, shapes, batchsizes, in_channels, out_channels, ksizes,
                strides):
traveller59's avatar
traveller59 committed
535
536
537
538
539
            device = torch.device(dev)
            num_points = [1000] * bs

            sparse_dict = generate_sparse_data(shape, num_points, IC)

540
541
542
543
            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
544
            features_dense = sparse_dict["features_dense"].astype(np.float32)
545
546
            filters = np.random.uniform(0, 1, size=[k, k, k, IC,
                                                    OC]).astype(np.float32)
traveller59's avatar
traveller59 committed
547
            indices_t = torch.from_numpy(indices).int().to(device)
548
549
            indices_scn_t = torch.from_numpy(
                indices[:, [1, 2, 3, 0]]).int().to(device)
traveller59's avatar
traveller59 committed
550
551
552
553
554
555
556
            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)
557
558
559
560
            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
561
562
            out_ref = net_ref(features_ref_t, indices_scn_t, bs)
            out = net(features_t, indices_t, bs)
563
564
            dout = np.random.uniform(-0.2, 0.2,
                                     out_ref.shape).astype(features.dtype)
traveller59's avatar
traveller59 committed
565
566
567
568
569
570
571
572
573
574
            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()
575
576
                dw_ref = layer_ref.weight.grad.detach().cpu().view(
                    *dw.shape).numpy()
traveller59's avatar
traveller59 committed
577
578
579
580
581
582
583
                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):
584
        np.random.seed(485)
585
        devices = ["cuda:0"]
traveller59's avatar
traveller59 committed
586
587
588
        shapes = [[19, 18, 17]]
        batchsizes = [1, 2]

yan.yan's avatar
yan.yan committed
589
590
        in_channels = [64]
        out_channels = [64]
traveller59's avatar
traveller59 committed
591
592
593
594
        ksizes = [2, 3]
        strides = [1, 2, 3]
        paddings = [0, 1]
        dilations = [1, 2, 3]
yan.yan's avatar
yan.yan committed
595
596
597
598
        ksizes = [2]
        strides = [2]
        paddings = [0]
        dilations = [1]
traveller59's avatar
traveller59 committed
599
600

        for dev, shape, bs, IC, OC, k, s, p, d in params_grid(
601
602
                devices, shapes, batchsizes, in_channels, out_channels, ksizes,
                strides, paddings, dilations):
traveller59's avatar
traveller59 committed
603
            if all([s > 1, d > 1]):
604
                continue  # don't support this.
traveller59's avatar
traveller59 committed
605
606
            device = torch.device(dev)
            num_points = [1000] * bs
yan.yan's avatar
yan.yan committed
607

traveller59's avatar
traveller59 committed
608
            # when data contains negative, sparse maxpool is not equal to dense maxpool.
609
610
611
612
613
614
615
616
617
            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
618
            features_dense = sparse_dict["features_dense"].astype(np.float32)
yan.yan's avatar
yan.yan committed
619
620
            filters = np.random.uniform(0, 1, size=[k, k, k, OC,
                                                    IC]).astype(np.float32)
traveller59's avatar
traveller59 committed
621
622
623
624
625
626
627
628
629
630
            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
631

traveller59's avatar
traveller59 committed
632
633
            outids = out.indices
            outfeatures = out.features
634
            outids_dev = outids.float()
635
            out_dense = out.dense(channels_first=False)
traveller59's avatar
traveller59 committed
636
            out = out_dense.permute(0, 4, 1, 2, 3).contiguous()
yan.yan's avatar
yan.yan committed
637
638
639
            out_np = out.detach().cpu().numpy()
            out_ref_np = out_ref.detach().cpu().numpy()
            self.assertAllClose(out_np, out_ref_np, atol=1e-4)
640
641
642

            dout_sparse = np.random.uniform(
                -0.2, 0.2, outfeatures.shape).astype(features.dtype)
traveller59's avatar
traveller59 committed
643
            dout_sparse_t = torch.from_numpy(dout_sparse).to(device)
644
645
            dout_t = scatter_nd(outids.long(), dout_sparse_t,
                                list(out_dense.shape))
traveller59's avatar
traveller59 committed
646
647
648
            dout_t = dout_t.permute(0, 4, 1, 2, 3).contiguous()
            out.backward(dout_t)
            out_ref.backward(dout_t)
649
650
            din_dense = features_dense_t.grad.detach().permute(0, 2, 3, 4,
                                                               1).contiguous()
traveller59's avatar
traveller59 committed
651
652
            din_sparse = gather_nd(din_dense, indices_t.long())
            din = features_t.grad.detach()
653

654
655
656
            din_np = din.cpu().numpy()
            din_sparse_np = din_sparse.cpu().numpy()
            self.assertAllClose(din_np, din_sparse_np, atol=1e-4)
657

traveller59's avatar
traveller59 committed
658

yanyan's avatar
yanyan committed
659
def main(algo=spconv.ConvAlgo.Native, dtype=torch.float32):
traveller59's avatar
traveller59 committed
660
661
    # function for develop.
    np.random.seed(484)
traveller59's avatar
traveller59 committed
662
    # devices = ["cuda:0"]
traveller59's avatar
traveller59 committed
663
    devices = ["cuda:0"]
yanyan's avatar
yanyan committed
664
    shapes = [[400, 400, 15]]
yanyan's avatar
yanyan committed
665
    batchsizes = [2]
traveller59's avatar
traveller59 committed
666

yanyan's avatar
yanyan committed
667
668
    in_channels = [19]
    out_channels = [17]
Yan Yan's avatar
Yan Yan committed
669
    ksizes = [(3, 3, 3)]
traveller59's avatar
traveller59 committed
670
671
672
673
674
    strides = [1]
    paddings = [0]
    dilations = [1]

    for dev, shape, bs, IC, OC, k, s, p, d in params_grid(
675
676
            devices, shapes, batchsizes, in_channels, out_channels, ksizes,
            strides, paddings, dilations):
traveller59's avatar
traveller59 committed
677
678
679
        if all([s > 1, d > 1]):
            continue
        device = torch.device(dev)
yanyan's avatar
yanyan committed
680
        num_points = [30000] * bs
traveller59's avatar
traveller59 committed
681
682
683

        sparse_dict = generate_sparse_data(shape, num_points, IC)

684
685
686
687
        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
688
689
        features_dense = sparse_dict["features_dense"].astype(np.float32)
        indices_t = torch.from_numpy(indices)
690
691
        filters = np.random.uniform(0, 1, size=[k[0], 1, 1, IC,
                                                OC]).astype(np.float32)
yanyan's avatar
yanyan committed
692
693
        indices_t = torch.from_numpy(indices).int().to(device).to(dtype)
        features_t = torch.from_numpy(features).to(device).to(dtype)
694

yanyan's avatar
yanyan committed
695
696
        features_dense_t = torch.from_numpy(features_dense).to(device).to(
            dtype)
yanyan's avatar
yanyan committed
697
        net = SparseConv3dTestTorch(1, 3, shape, IC, OC, k, s, p, d,
yanyan's avatar
yanyan committed
698
                                    algo=algo).to(device).to(dtype)
699
        net_ref = Conv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
yanyan's avatar
yanyan committed
700
701
                                  d).to(device).to(dtype)
        filters_t = torch.from_numpy(filters).to(device).to(dtype)
702
703
        net_ref.net[0].weight[:] = filters_t.permute(4, 3, 0, 1,
                                                     2).contiguous()
traveller59's avatar
traveller59 committed
704
705
706
        net.net[0].weight[:] = filters_t
        out_ref = net_ref(features_dense_t)
        times = []
tusimple's avatar
tusimple committed
707
        for i in range(10):
traveller59's avatar
traveller59 committed
708
709
710
711
            t = time.time()
            out = net(features_t, indices_t, bs)
            torch.cuda.synchronize()
            times.append(time.time() - t)
712
        # print((net.grid == -1).float().sum(), net.grid.numel())
713
        # print("spconv time", time.time() - t)
traveller59's avatar
traveller59 committed
714
        print("spconv time", np.mean(times[2:]))
traveller59's avatar
traveller59 committed
715
716
717
        out = net(features_t, indices_t, bs)
        # print(out.indices)
        out = out.dense()
Yan Yan's avatar
Yan Yan committed
718
719
        out_numpy = out.detach().cpu().numpy()

720
721
722
        print(
            np.linalg.norm(out.detach().cpu().numpy() -
                           out_ref.detach().cpu().numpy()))
yanyan's avatar
yanyan committed
723
724
        print(out_numpy.min(), out_numpy.max(), out_numpy.mean(),
              out_numpy.sum())
traveller59's avatar
traveller59 committed
725
726


yanyan's avatar
yanyan committed
727
def main_subm(algo, dtype=torch.float32):
728
729
    # function for develop.
    np.random.seed(484)
Yan Yan's avatar
Yan Yan committed
730
    torch.manual_seed(50051)
731
732
    # devices = ["cuda:0"]
    devices = ["cuda:0"]
yanyan's avatar
yanyan committed
733
    shapes = [[400, 400, 15]]
734
735
736
737
738
739
740
741
742
743
744
745
746
747
    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
748
        num_points = [120000] * bs
749
750
751
752
753
754
755
756
757
758
759

        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
760
761
        indices_t = torch.from_numpy(indices).int().to(device).to(dtype)
        features_t = torch.from_numpy(features).to(device).to(dtype)
762

yanyan's avatar
yanyan committed
763
764
        features_dense_t = torch.from_numpy(features_dense).to(device).to(
            dtype)
yanyan's avatar
yanyan committed
765
        net = SubMConv3dTestTorch(1, 3, shape, IC, OC, k, s, p, d,
yanyan's avatar
yanyan committed
766
                                  algo=algo).to(device).to(dtype)
767
        net_ref = Conv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
yanyan's avatar
yanyan committed
768
769
                                  d).to(device).to(dtype)
        filters_t = torch.from_numpy(filters).to(device).to(dtype)
770
771
772
773
774
        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
775
        for i in range(20):
776
777
778
779
780
781
782
783
784
785
            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
786
        out_numpy = out.detach().cpu().numpy()
yanyan's avatar
yanyan committed
787
788
789
        # print(
        #     np.linalg.norm(out.detach().cpu().numpy() -
        #                    out_ref.detach().cpu().numpy()))
yanyan's avatar
yanyan committed
790
791
        print(out_numpy.min(), out_numpy.max(), out_numpy.mean(),
              out_numpy.sum())
Yan Yan's avatar
Yan Yan committed
792
    return out_numpy
793

yanyan's avatar
yanyan committed
794

traveller59's avatar
traveller59 committed
795
if __name__ == '__main__':
yanyan's avatar
yanyan committed
796
797
    # main_subm(algo=spconv.ConvAlgo.SparseConvNet, dtype=torch.float32)
    # main(algo=spconv.ConvAlgo.SparseConvNet, dtype=torch.float32)
Yan Yan's avatar
Yan Yan committed
798
    # TestCase().assertAllClose(out_my, out_ref)
traveller59's avatar
traveller59 committed
799
    # unittest.main()
yanyan's avatar
yanyan committed
800
    TestSpConv().testSpConv3d()