test_conv.py 30 KB
Newer Older
traveller59's avatar
traveller59 committed
1
# Copyright 2019 Yan Yan
2
#
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
6
#
traveller59's avatar
traveller59 committed
7
#     http://www.apache.org/licenses/LICENSE-2.0
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
23
24
25
26
27
from torch import nn

import spconv
from spconv.test_utils import TestCase, generate_sparse_data, params_grid

# import sparseconvnet as scn

traveller59's avatar
traveller59 committed
28
29

class SparseConv3dTestTorch(nn.Module):
30
31
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride, padding, dilation):
traveller59's avatar
traveller59 committed
32
        super().__init__()
Yan Yan's avatar
Yan Yan committed
33
34
        algo = spconv.ConvAlgo.BatchGemm

35
36
37
38
39
40
41
42
        layers = [
            spconv.SparseConv3d(in_channels,
                                out_channels,
                                kernel_size,
                                stride,
                                padding=padding,
                                dilation=dilation,
                                bias=False,
Yan Yan's avatar
Yan Yan committed
43
44
                                use_hash=False,
                                algo=algo)
45
        ]
traveller59's avatar
traveller59 committed
46
        for i in range(1, num_layers):
47
48
49
50
51
52
53
            layers.append(
                spconv.SparseConv3d(out_channels,
                                    out_channels,
                                    kernel_size,
                                    stride,
                                    padding=padding,
                                    dilation=dilation,
54
                                    bias=False,
Yan Yan's avatar
Yan Yan committed
55
56
                                    use_hash=False,
                                    algo=algo))
57
        self.net = spconv.SparseSequential(*layers, )
traveller59's avatar
traveller59 committed
58
59
60
61
62
63
        # 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()
64
65
66
67
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size,
                                    self.grid)
        return self.net(x)  # .dense()

traveller59's avatar
traveller59 committed
68
69

class SubMConv3dTestTorch(nn.Module):
70
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
Yan Yan's avatar
Yan Yan committed
71
                 kernel_size, stride, padding, dilation, algo=spconv.ConvAlgo.Native):
traveller59's avatar
traveller59 committed
72
        super().__init__()
73
74
75
76
77
78
79
        layers = [
            spconv.SubMConv3d(in_channels,
                              out_channels,
                              kernel_size,
                              stride,
                              padding=padding,
                              dilation=dilation,
Yan Yan's avatar
Yan Yan committed
80
81
                              bias=False,
                              algo=algo)
82
        ]
traveller59's avatar
traveller59 committed
83
        for i in range(1, num_layers):
84
85
86
87
88
89
90
            layers.append(
                spconv.SubMConv3d(out_channels,
                                  out_channels,
                                  kernel_size,
                                  stride,
                                  padding=padding,
                                  dilation=dilation,
Yan Yan's avatar
Yan Yan committed
91
92
                                  bias=False,
                                algo=algo))
93
        self.net = spconv.SparseSequential(*layers, )
94
95
        # self.grid = torch.full([3, *shape], -1, dtype=torch.int32).cuda()
        self.grid = None
traveller59's avatar
traveller59 committed
96
97
98
        self.shape = shape

    def forward(self, features, coors, batch_size):
Yan Yan's avatar
Yan Yan committed
99
        coors = coors.int()# .cpu()
100
101
102
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size,
                                    self.grid)
        return self.net(x)  # .dense()
traveller59's avatar
traveller59 committed
103
104
105


class Conv3dTestTorch(nn.Module):
106
107
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride, padding, dilation):
traveller59's avatar
traveller59 committed
108
        super().__init__()
109
110
111
112
113
114
115
116
117
        layers = [
            nn.Conv3d(in_channels,
                      out_channels,
                      kernel_size,
                      stride,
                      padding=padding,
                      dilation=dilation,
                      bias=False)
        ]
traveller59's avatar
traveller59 committed
118
        for i in range(1, num_layers):
119
120
121
122
123
124
125
126
127
            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
128
129
130
        self.shape = shape

    def forward(self, x):
131
        return self.net(x)  # .dense()
traveller59's avatar
traveller59 committed
132
133
134


class SparseDeConv3dTestTorch(nn.Module):
135
136
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride, padding, dilation):
traveller59's avatar
traveller59 committed
137
        super().__init__()
138
139
140
141
142
143
144
145
146
        layers = [
            spconv.SparseConvTranspose3d(in_channels,
                                         out_channels,
                                         kernel_size,
                                         stride,
                                         padding=padding,
                                         dilation=dilation,
                                         bias=False)
        ]
traveller59's avatar
traveller59 committed
147
        for i in range(1, num_layers):
148
149
150
151
152
153
154
155
156
            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
157
158
159
160
        self.shape = shape

    def forward(self, features, coors, batch_size):
        coors = coors.int()
161
162
163
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size)
        return self.net(x)  # .dense()

traveller59's avatar
traveller59 committed
164
165

class DeConv3dTestTorch(nn.Module):
166
167
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride, padding, dilation):
traveller59's avatar
traveller59 committed
168
        super().__init__()
169
170
171
172
173
174
175
176
177
        layers = [
            nn.ConvTranspose3d(in_channels,
                               out_channels,
                               kernel_size,
                               stride,
                               padding=padding,
                               dilation=dilation,
                               bias=False)
        ]
traveller59's avatar
traveller59 committed
178
        for i in range(1, num_layers):
179
180
181
182
183
184
185
186
187
            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
188
189
190
        self.shape = shape

    def forward(self, x):
191
        return self.net(x)  # .dense()
traveller59's avatar
traveller59 committed
192
193
194


class SparseMaxPoolTestTorch(nn.Module):
195
196
    def __init__(self, num_layers, ndim, shape, kernel_size, stride, padding,
                 dilation):
traveller59's avatar
traveller59 committed
197
        super().__init__()
198
199
200
        layers = [
            spconv.SparseMaxPool3d(kernel_size, stride, padding, dilation)
        ]
traveller59's avatar
traveller59 committed
201
        for i in range(1, num_layers):
202
203
204
            layers.append(
                spconv.SparseMaxPool3d(kernel_size, stride, padding, dilation))
        self.net = spconv.SparseSequential(*layers, )
traveller59's avatar
traveller59 committed
205
206
207
208
        self.shape = shape

    def forward(self, features, coors, batch_size):
        coors = coors.int()
209
210
211
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size)
        return self.net(x)  # .dense()

traveller59's avatar
traveller59 committed
212
213

class MaxPool3dTestTorch(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
        layers = [nn.MaxPool3d(kernel_size, stride, padding, dilation)]
traveller59's avatar
traveller59 committed
218
        for i in range(1, num_layers):
219
220
            layers.append(nn.MaxPool3d(kernel_size, stride, padding, dilation))
        self.net = nn.Sequential(*layers, )
traveller59's avatar
traveller59 committed
221
222
223
        self.shape = shape

    def forward(self, x):
224
        return self.net(x)  # .dense()
traveller59's avatar
traveller59 committed
225
226
227


class SubmanifoldConvTestTorch(nn.Module):
228
229
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride):
traveller59's avatar
traveller59 committed
230
        super().__init__()
231
232
233
234
235
236
237
        layers = [
            spconv.SubMConv3d(in_channels,
                              out_channels,
                              kernel_size,
                              bias=False,
                              indice_key="subm0")
        ]
traveller59's avatar
traveller59 committed
238
        for i in range(1, num_layers):
239
240
241
242
243
244
            layers.append(
                spconv.SubMConv3d(out_channels,
                                  out_channels,
                                  kernel_size,
                                  bias=False))
        self.net = nn.Sequential(*layers, )
traveller59's avatar
traveller59 committed
245
246
247
248
        self.shape = shape

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

252

traveller59's avatar
traveller59 committed
253
class SCNCoupleDeConvTest(nn.Module):
254
255
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride):
traveller59's avatar
traveller59 committed
256
257
258
        super().__init__()
        self.scn_input = scn.InputLayer(ndim, shape, mode=0)
        self.net = nn.Sequential(
259
260
261
262
263
264
265
266
267
268
269
270
            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
271
272
273
274
275
276
277
278
            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)

279

traveller59's avatar
traveller59 committed
280
class SparseCoupleDeConvTest(nn.Module):
281
282
    def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
                 kernel_size, stride):
traveller59's avatar
traveller59 committed
283
284
        super().__init__()
        self.net = spconv.SparseSequential(
285
286
287
288
289
290
291
292
293
294
295
            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
296
297
298
299
300
301
        )
        self.todense = spconv.ToDense()
        self.shape = shape

    def forward(self, features, coors, batch_size):
        coors = coors.int()
302
303
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size)
        return self.todense(self.net(x))  # .dense()
traveller59's avatar
traveller59 committed
304
305
306
307
308


def gather_nd(params, indices):
    # this function has a limit that MAX_ADVINDEX_CALC_DIMS=5
    ndim = indices.shape[-1]
309
310
    output_shape = list(indices.shape[:-1]) + list(
        params.shape[indices.shape[-1]:])
traveller59's avatar
traveller59 committed
311
312
313
314
315
    flatted_indices = indices.view(-1, ndim)
    slices = [flatted_indices[:, i] for i in range(ndim)]
    slices += [Ellipsis]
    return params[slices].view(*output_shape)

316

traveller59's avatar
traveller59 committed
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
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
Yan Yan committed
336
        devices = ["cuda:0"]
traveller59's avatar
traveller59 committed
337
338
339
340
341
342
343
344
345
346
347
        shapes = [[19, 18, 17]]
        batchsizes = [1, 2]

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

        for dev, shape, bs, IC, OC, k, s, p, d in params_grid(
348
349
                devices, shapes, batchsizes, in_channels, out_channels, ksizes,
                strides, paddings, dilations):
traveller59's avatar
traveller59 committed
350
            if all([s > 1, d > 1]):
351
                continue  # don't support this.
traveller59's avatar
traveller59 committed
352
353
354
355
356
            device = torch.device(dev)
            num_points = [1000] * bs

            sparse_dict = generate_sparse_data(shape, num_points, IC)

357
358
359
360
            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
361
            features_dense = sparse_dict["features_dense"].astype(np.float32)
362
363
            filters = np.random.uniform(0, 1, size=[k, k, k, IC,
                                                    OC]).astype(np.float32)
traveller59's avatar
traveller59 committed
364
365
366
367
368
            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
369
370
371
372
            net = SparseConv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
                                        d).to(device)
            net_ref = Conv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
                                      d).to(device)
traveller59's avatar
traveller59 committed
373
            filters_t = torch.from_numpy(filters).to(device)
374
375
            net_ref.net[0].weight.data[:] = filters_t.permute(4, 3, 0, 1,
                                                              2).contiguous()
traveller59's avatar
traveller59 committed
376
377
            net.net[0].weight.data[:] = filters_t
            out_ref = net_ref(features_dense_t)
378
            out = net(features_t, indices_t, bs).dense()
379
380
            dout = np.random.uniform(-0.2, 0.2,
                                     out_ref.shape).astype(features.dtype)
traveller59's avatar
traveller59 committed
381
382
383
            dout_t = torch.from_numpy(dout).to(device)
            out.backward(dout_t)
            out_ref.backward(dout_t)
384
385
            din_dense = features_dense_t.grad.detach().permute(0, 2, 3, 4,
                                                               1).contiguous()
traveller59's avatar
traveller59 committed
386
387
388
389
390
391
392
393
394
395
396
397
398
399
            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()
                dw = dw.transpose(4, 3, 0, 1, 2)
                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)
400

traveller59's avatar
traveller59 committed
401
402
    def testSpDeConv3d(self):
        np.random.seed(484)
Yan Yan's avatar
Yan Yan committed
403
        devices = ["cuda:0"]
traveller59's avatar
traveller59 committed
404
405
406
407
408
409
410
411
412
413
414
        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]

        for dev, shape, bs, IC, OC, k, s, p, d in params_grid(
415
416
                devices, shapes, batchsizes, in_channels, out_channels, ksizes,
                strides, paddings, dilations):
traveller59's avatar
traveller59 committed
417
            if all([s > 1, d > 1]):
418
                continue  # don't support this.
traveller59's avatar
traveller59 committed
419
420
421
422
423
            device = torch.device(dev)
            num_points = [1000] * bs

            sparse_dict = generate_sparse_data(shape, num_points, IC)

424
425
426
427
            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
428
            features_dense = sparse_dict["features_dense"].astype(np.float32)
429
430
            filters = np.random.uniform(0, 1, size=[k, k, k, IC,
                                                    OC]).astype(np.float32)
traveller59's avatar
traveller59 committed
431
432
433
434
435
            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
436
437
438
439
            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
440
            filters_t = torch.from_numpy(filters).to(device)
441
442
            net_ref.net[0].weight.data[:] = filters_t.permute(3, 4, 0, 1,
                                                              2).contiguous()
traveller59's avatar
traveller59 committed
443
444
            net.net[0].weight.data[:] = filters_t
            out_ref = net_ref(features_dense_t)
445
            out = net(features_t, indices_t, bs).dense()
446
447
            dout = np.random.uniform(-0.2, 0.2,
                                     out_ref.shape).astype(features.dtype)
traveller59's avatar
traveller59 committed
448
449
450
            dout_t = torch.from_numpy(dout).to(device)
            out.backward(dout_t)
            out_ref.backward(dout_t)
451
452
            din_dense = features_dense_t.grad.detach().permute(0, 2, 3, 4,
                                                               1).contiguous()
traveller59's avatar
traveller59 committed
453
454
455
456
457
458
459
460
461
462
463
464
465
466
            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()
                dw = dw.transpose(3, 4, 0, 1, 2)
                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)
467

traveller59's avatar
traveller59 committed
468
469
470
471
472
473
474
475
476
477
478
479
480
481
    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(
482
483
                devices, shapes, batchsizes, in_channels, out_channels, ksizes,
                strides):
traveller59's avatar
traveller59 committed
484
485
486
487
488
            device = torch.device(dev)
            num_points = [1000] * bs

            sparse_dict = generate_sparse_data(shape, num_points, IC)

489
490
491
492
            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
493
            features_dense = sparse_dict["features_dense"].astype(np.float32)
494
495
            filters = np.random.uniform(0, 1, size=[k, k, k, IC,
                                                    OC]).astype(np.float32)
traveller59's avatar
traveller59 committed
496
            indices_t = torch.from_numpy(indices).int().to(device)
497
498
            indices_scn_t = torch.from_numpy(
                indices[:, [1, 2, 3, 0]]).int().to(device)
traveller59's avatar
traveller59 committed
499
500
501
502
503
504
505
            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)
506
507
508
509
            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
510
511
            out_ref = net_ref(features_ref_t, indices_scn_t, bs)
            out = net(features_t, indices_t, bs)
512
513
            dout = np.random.uniform(-0.2, 0.2,
                                     out_ref.shape).astype(features.dtype)
traveller59's avatar
traveller59 committed
514
515
516
517
518
519
520
521
522
523
            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()
524
525
                dw_ref = layer_ref.weight.grad.detach().cpu().view(
                    *dw.shape).numpy()
traveller59's avatar
traveller59 committed
526
527
528
529
530
531
532
                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):
533
        np.random.seed(485)
534
        devices = ["cuda:0"]
traveller59's avatar
traveller59 committed
535
536
537
        shapes = [[19, 18, 17]]
        batchsizes = [1, 2]

538
539
        in_channels = [62]
        out_channels = [62]
traveller59's avatar
traveller59 committed
540
541
542
543
544
545
        ksizes = [2, 3]
        strides = [1, 2, 3]
        paddings = [0, 1]
        dilations = [1, 2, 3]

        for dev, shape, bs, IC, OC, k, s, p, d in params_grid(
546
547
                devices, shapes, batchsizes, in_channels, out_channels, ksizes,
                strides, paddings, dilations):
traveller59's avatar
traveller59 committed
548
            if all([s > 1, d > 1]):
549
                continue  # don't support this.
traveller59's avatar
traveller59 committed
550
551
552
            device = torch.device(dev)
            num_points = [1000] * bs
            # when data contains negative, sparse maxpool is not equal to dense maxpool.
553
554
555
556
557
558
559
560
561
            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
562
            features_dense = sparse_dict["features_dense"].astype(np.float32)
563
564
            filters = np.random.uniform(0, 1, size=[k, k, k, IC,
                                                    OC]).astype(np.float32)
traveller59's avatar
traveller59 committed
565
566
567
568
569
570
571
572
573
574
575
576
            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)
            outids = out.indices
            outfeatures = out.features
577
            outids_dev = outids.float()
578
            out_dense = out.dense(channels_first=False)
traveller59's avatar
traveller59 committed
579
            out = out_dense.permute(0, 4, 1, 2, 3).contiguous()
580
581
582

            dout_sparse = np.random.uniform(
                -0.2, 0.2, outfeatures.shape).astype(features.dtype)
traveller59's avatar
traveller59 committed
583
            dout_sparse_t = torch.from_numpy(dout_sparse).to(device)
584
585
            dout_t = scatter_nd(outids.long(), dout_sparse_t,
                                list(out_dense.shape))
traveller59's avatar
traveller59 committed
586
587
588
            dout_t = dout_t.permute(0, 4, 1, 2, 3).contiguous()
            out.backward(dout_t)
            out_ref.backward(dout_t)
589
590
            din_dense = features_dense_t.grad.detach().permute(0, 2, 3, 4,
                                                               1).contiguous()
traveller59's avatar
traveller59 committed
591
592
            din_sparse = gather_nd(din_dense, indices_t.long())
            din = features_t.grad.detach()
593

traveller59's avatar
traveller59 committed
594
595
596
            out_np = out.detach().cpu().numpy()
            out_ref_np = out_ref.detach().cpu().numpy()
            self.assertAllClose(out_np, out_ref_np, atol=1e-4)
597
598
599
            din_np = din.cpu().numpy()
            din_sparse_np = din_sparse.cpu().numpy()
            self.assertAllClose(din_np, din_sparse_np, atol=1e-4)
600

traveller59's avatar
traveller59 committed
601
602
603
604

def main():
    # function for develop.
    np.random.seed(484)
traveller59's avatar
traveller59 committed
605
    # devices = ["cuda:0"]
traveller59's avatar
traveller59 committed
606
607
    devices = ["cuda:0"]
    shapes = [[50, 30, 30]]
traveller59's avatar
traveller59 committed
608
    batchsizes = [2]
traveller59's avatar
traveller59 committed
609

Yan Yan's avatar
Yan Yan committed
610
611
612
    in_channels = [32]
    out_channels = [64]
    ksizes = [(3, 3, 3)]
traveller59's avatar
traveller59 committed
613
614
615
616
617
    strides = [1]
    paddings = [0]
    dilations = [1]

    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
621
622
        if all([s > 1, d > 1]):
            continue
        device = torch.device(dev)
tusimple's avatar
tusimple committed
623
        num_points = [500] * bs
traveller59's avatar
traveller59 committed
624
625
626

        sparse_dict = generate_sparse_data(shape, num_points, IC)

627
628
629
630
        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
631
632
        features_dense = sparse_dict["features_dense"].astype(np.float32)
        indices_t = torch.from_numpy(indices)
633
634
        filters = np.random.uniform(0, 1, size=[k[0], 1, 1, IC,
                                                OC]).astype(np.float32)
traveller59's avatar
traveller59 committed
635
636
        indices_t = torch.from_numpy(indices).int().to(device).float()
        features_t = torch.from_numpy(features).to(device).float()
637

traveller59's avatar
traveller59 committed
638
        features_dense_t = torch.from_numpy(features_dense).to(device).float()
639
640
641
642
        net = SparseConv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
                                    d).to(device).float()
        net_ref = Conv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
                                  d).to(device).float()
traveller59's avatar
traveller59 committed
643
        filters_t = torch.from_numpy(filters).to(device).float()
644
645
        net_ref.net[0].weight[:] = filters_t.permute(4, 3, 0, 1,
                                                     2).contiguous()
traveller59's avatar
traveller59 committed
646
647
648
        net.net[0].weight[:] = filters_t
        out_ref = net_ref(features_dense_t)
        times = []
tusimple's avatar
tusimple committed
649
        for i in range(10):
traveller59's avatar
traveller59 committed
650
651
652
653
            t = time.time()
            out = net(features_t, indices_t, bs)
            torch.cuda.synchronize()
            times.append(time.time() - t)
654
        # print((net.grid == -1).float().sum(), net.grid.numel())
655
        # print("spconv time", time.time() - t)
traveller59's avatar
traveller59 committed
656
        print("spconv time", np.mean(times[2:]))
traveller59's avatar
traveller59 committed
657
658
659
        out = net(features_t, indices_t, bs)
        # print(out.indices)
        out = out.dense()
Yan Yan's avatar
Yan Yan committed
660
661
        out_numpy = out.detach().cpu().numpy()

662
663
664
        print(
            np.linalg.norm(out.detach().cpu().numpy() -
                           out_ref.detach().cpu().numpy()))
Yan Yan's avatar
Yan Yan committed
665
        print(out_numpy.min(), out_numpy.max(), out_numpy.mean(), out_numpy.sum())
traveller59's avatar
traveller59 committed
666
667


Yan Yan's avatar
Yan Yan committed
668
def main_subm(algo):
669
670
    # function for develop.
    np.random.seed(484)
Yan Yan's avatar
Yan Yan committed
671
    torch.manual_seed(50051)
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
    # devices = ["cuda:0"]
    devices = ["cuda:0"]
    shapes = [[50, 30, 30]]
    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)
Yan Yan's avatar
Yan Yan committed
689
        num_points = [1000] * bs
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705

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

        features_dense_t = torch.from_numpy(features_dense).to(device).float()
        net = SubMConv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
Yan Yan's avatar
Yan Yan committed
706
                                  d, algo=algo).to(device).float()
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
        net_ref = Conv3dTestTorch(1, 3, shape, IC, OC, k, s, p,
                                  d).to(device).float()
        filters_t = torch.from_numpy(filters).to(device).float()
        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 = []
        for i in range(100):
            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
726
        out_numpy = out.detach().cpu().numpy()
727
728
729
        print(
            np.linalg.norm(out.detach().cpu().numpy() -
                           out_ref.detach().cpu().numpy()))
Yan Yan's avatar
Yan Yan committed
730
731
        print(out_numpy.min(), out_numpy.max(), out_numpy.mean(), out_numpy.sum())
    return out_numpy
732

traveller59's avatar
traveller59 committed
733
if __name__ == '__main__':
Yan Yan's avatar
Yan Yan committed
734
735
736
    # out_my = main_subm(algo=spconv.ConvAlgo.BatchGemm)
    # out_ref = main_subm(algo=spconv.ConvAlgo.Native)
    # TestCase().assertAllClose(out_my, out_ref)
traveller59's avatar
traveller59 committed
737
    # unittest.main()
Yan Yan's avatar
Yan Yan committed
738
    TestSpConv().testSpConv3d()