benchmark.py 5.7 KB
Newer Older
yanyan's avatar
yanyan committed
1
2
3
4
5
import time
from pathlib import Path

import numpy as np
import torch
yanyan's avatar
yanyan committed
6
from torch import nn
yanyan's avatar
yanyan committed
7
8
9
10

import spconv
from spconv.utils import VoxelGeneratorV2

yanyan's avatar
yanyan committed
11
12

def waymo_data(batch_size=1):
yanyan's avatar
yanyan committed
13
14
    gen = VoxelGeneratorV2([0.1, 0.1, 0.1], [-80, -80, -2, 80, 80, 6], 1,
                           150000)
yanyan's avatar
yanyan committed
15
16
17
18
19
20
21
22
23
    data = np.load(Path(__file__).parent / "data" / "benchmark-pc.npz")
    pc = data["pc"]
    data = gen.generate(pc)
    voxels = data["voxels"].reshape(-1, 3)
    coors = data["coordinates"]
    N = coors.shape[0]
    coors = np.concatenate([np.full([N, 1], 0, coors.dtype), coors], axis=1)
    return voxels, coors, gen.grid_size

yanyan's avatar
yanyan committed
24

yanyan's avatar
yanyan committed
25
class Net(nn.Module):
yanyan's avatar
yanyan committed
26
    def __init__(self, shape, algo):
yanyan's avatar
yanyan committed
27
28
        super().__init__()
        self.net = spconv.SparseSequential(
yanyan's avatar
yanyan committed
29
30
31
32
33
34
35
36
            spconv.SubMConv3d(3, 64, 3, bias=False, indice_key="c0",
                              algo=algo),
            spconv.SubMConv3d(64,
                              64,
                              3,
                              bias=False,
                              indice_key="c0",
                              algo=algo),
yanyan's avatar
yanyan committed
37
38
39
            # nn.BatchNorm1d(32),
            # nn.ReLU(),
            spconv.SparseMaxPool3d(2, 2),
yanyan's avatar
yanyan committed
40
41
42
43
44
45
46
47
48
49
50
51
            spconv.SubMConv3d(64,
                              96,
                              3,
                              bias=False,
                              indice_key="c1",
                              algo=algo),
            spconv.SubMConv3d(96,
                              96,
                              3,
                              bias=False,
                              indice_key="c1",
                              algo=algo),
yanyan's avatar
yanyan committed
52
53
54
            # nn.BatchNorm1d(64),
            # nn.ReLU(),
            spconv.SparseMaxPool3d(2, 2),
yanyan's avatar
yanyan committed
55
56
57
58
59
60
61
62
63
64
65
66
            spconv.SubMConv3d(96,
                              128,
                              3,
                              bias=False,
                              indice_key="c2",
                              algo=algo),
            spconv.SubMConv3d(128,
                              128,
                              3,
                              bias=False,
                              indice_key="c2",
                              algo=algo),
yanyan's avatar
yanyan committed
67
68
69
            # nn.BatchNorm1d(128),
            # nn.ReLU(),
            spconv.SparseMaxPool3d(2, 2),
yanyan's avatar
yanyan committed
70
71
72
73
74
75
76
77
78
79
80
81
            spconv.SubMConv3d(128,
                              160,
                              3,
                              bias=False,
                              indice_key="c3",
                              algo=algo),
            spconv.SubMConv3d(160,
                              160,
                              3,
                              bias=False,
                              indice_key="c3",
                              algo=algo),
yanyan's avatar
yanyan committed
82
83
84
            # nn.BatchNorm1d(128),
            # nn.ReLU(),
            spconv.SparseMaxPool3d(2, 2),
yanyan's avatar
yanyan committed
85
86
87
88
89
90
91
92
93
94
95
96
            spconv.SubMConv3d(160,
                              192,
                              3,
                              bias=False,
                              indice_key="c4",
                              algo=algo),
            spconv.SubMConv3d(192,
                              192,
                              3,
                              bias=False,
                              indice_key="c4",
                              algo=algo),
yanyan's avatar
yanyan committed
97
98
99
            # nn.BatchNorm1d(128),
            # nn.ReLU(),
            spconv.SparseMaxPool3d(2, 2),
yanyan's avatar
yanyan committed
100
101
102
103
104
105
106
107
108
109
110
111
            spconv.SubMConv3d(192,
                              224,
                              3,
                              bias=False,
                              indice_key="c5",
                              algo=algo),
            spconv.SubMConv3d(224,
                              224,
                              3,
                              bias=False,
                              indice_key="c5",
                              algo=algo),
yanyan's avatar
yanyan committed
112
113
114
            # nn.BatchNorm1d(128),
            # nn.ReLU(),
            spconv.SparseMaxPool3d(2, 2),
yanyan's avatar
yanyan committed
115
116
117
118
119
120
121
122
123
124
125
126
            spconv.SubMConv3d(224,
                              256,
                              3,
                              bias=False,
                              indice_key="c6",
                              algo=algo),
            spconv.SubMConv3d(256,
                              256,
                              3,
                              bias=False,
                              indice_key="c6",
                              algo=algo),
yanyan's avatar
yanyan committed
127
128
        )
        max_batch_size = 1
yanyan's avatar
yanyan committed
129
        # grid (dense map) is used for indice generation. use pre-allocated grid can run faster.
yanyan's avatar
yanyan committed
130
131
        self.grid = torch.full([max_batch_size, *shape], -1,
                               dtype=torch.int32).cuda()
yanyan's avatar
yanyan committed
132
133
134
135
136
137
138
139
140
141
142
143
        # self.grid = None
        self.shape = shape

    def forward(self, features, coors, batch_size):
        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size,
                                    self.grid)
        return self.net(x)


def main():
    voxels, coors, spatial_shape = waymo_data()
    voxels_th = torch.from_numpy(voxels).cuda().float()
yanyan's avatar
yanyan committed
144
    coors_th = torch.from_numpy(coors).cuda().int()
yanyan's avatar
yanyan committed
145
146
    algo = spconv.ConvAlgo.Native
    net = Net(spatial_shape[::-1], algo).cuda().eval().float()
yanyan's avatar
yanyan committed
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
    print(coors_th.shape)
    out = net(voxels_th, coors_th, 1)
    print(out.spatial_shape)
    times = []
    with torch.no_grad():
        for i in range(20):
            torch.cuda.synchronize()
            t = time.time()
            out = net(voxels_th, coors_th, 1)
            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:]))

yanyan's avatar
yanyan committed
162

yanyan's avatar
yanyan committed
163
if __name__ == "__main__":
yanyan's avatar
yanyan committed
164
    main()