Point-Voxel CNN for Efficient 3D Deep Learning

Introduction

3D deep learning has received increased attention thanks to its wide applications: e.g., AR/VR and autonomous driving. These applications need to interact with people in real time and therefore require low latency. However, edge devices (such as mobile phones and VR headsets) are tightly constrained by hardware resources and battery. Therefore, it is important to design efficient and fast 3D deep learning models for real-time applications on the edge.

Collected by the LiDAR sensors, 3D data usually comes in the format of point clouds. Conventionally, researchers rasterize the point cloud into voxel grids and process them using 3D volumetric convolutions . With low resolutions, there will be information loss during voxelization: multiple points will be merged together if they lie in the same grid. Therefore, a high-resolution representation is needed to preserve the fine details in the input data. However, the computational cost and memory requirement both increase cubically with voxel resolution. Thus, it is infeasible to train a voxel-based model with high-resolution inputs: e.g., 3D-UNet requires more than 10 GB of GPU memory on 64$\times$64$\times$64 inputs with batch size of 16, and the large memory footprint makes it rather difficult to scale beyond this resolution.

Recently, another stream of models attempt to directly process the input point clouds . These point-based models require much lower GPU memory than voxel-based models thanks to the sparse representation. However, they neglect the fact that the random memory access is also very inefficient. As the points are scattered over the entire 3D space in an irregular manner, processing them introduces random memory accesses. Most point-based models mimic the 3D volumetric convolution: they extract the feature of each point by aggregating its neighboring features. However, neighbors are not stored contiguously in the point representation; therefore, indexing them requires the costly nearest neighbor search. To trade space for time, previous methods replicate the entire point cloud for each center point in the nearest neighbor search, and the memory cost will then be $\mathcal{O}(n^{2})$, where $n$ is the number of input points. Another overhead is introduced by the dynamic kernel computation. Since the relative positions of neighbors are not fixed, these point-based models have to generate the convolution kernels dynamically based on different offsets.

Designing efficient 3D neural network models needs to take the hardware into account. Compared with arithmetic operations, memory operations are particularly expensive: they consume two orders of magnitude higher energy, having two orders of magnitude lower bandwidth (Figure 1(a)). Another aspect is the memory access pattern: the random access will introduce memory bank conflicts and decrease the throughput (Figure 1(b)). From the hardware perspective, conventional 3D models are inefficient due to large memory footprint and random memory access.

This paper provides a novel perspective to overcome these challenges. We propose Point-Voxel CNN (PVCNN) that represents the 3D input data as point clouds to take advantage of the sparsity to reduce the memory footprint, and leverages the voxel-based convolution to obtain the contiguous memory access pattern. Extensive experiments on multiple tasks demonstrate that PVCNN outperforms the voxel-based baseline with 10$\times$ lower memory consumption. It also achieves 7$\times$ measured speedup on average compared with the state-of-the-art point-based models.

Related Work

Extensive attention has been paid to hardware-efficient deep learning for real-world applications. For instance, researchers have proposed to reduce the memory access cost by pruning and quantizing the models or directly designing the compact models . However, all these approaches are general-purpose and are suitable for arbitrary neural networks. In this paper, we instead design our efficient primitive based on some domain-specific properties: e.g., 3D point clouds are highly sparse and spatially structured.

Voxel-Based 3D Models.

Conventionally, researchers relied on the volumetric representation to process the 3D data . For instance, Maturana et al. proposed the vanilla volumetric CNN; Qi et al. extended 2D CNNs to 3D and systematically analyzed the relationship between 3D CNNs and multi-view CNNs; Wang et al. incoporated the octree into volumetric CNNs to reduce the memory consumption. Recent studies suggest that the volumetric representation can also be used in 3D shape segmentation and 3D object detection .

Point-Based 3D Models.

PointNet takes advantage of the symmetric function to process the unordered point sets in 3D. Later research proposed to stack PointNets hierarchically to model neighborhood information and increase model capacity. Instead of stacking PointNets as basic blocks, another type of methods abstract away the symmetric function using dynamically generated convolution kernels or learned neighborhood permutation function. Other research, such as SPLATNet which naturally extends the idea of 2D image SPLAT to 3D, and SONet which uses the self-organization mechanism with the theoretical guarantee of invariance to point order, also shows great potential in general-purpose 3D modeling with point clouds as input.

Special-Purpose 3D Models.

There are also 3D models tailored for specific tasks. For instance, SegCloud , SGPN , SPGraph , ParamConv , SSCN and RSNet are specialized in 3D semantic/instance segmentation. As for 3D object detection, F-PointNet is based on the RGB detector and point-based regional proposal networks; PointRCNN follows the similar idea while abstracting away the RGB detector; PointPillars and SECOND focus on the efficiency.

Motivation

3D data can be represented in the format of $\bm{x}=\{\bm{x}_{k}\}=\{(\bm{p}_{k},\bm{f}_{k})\}$, where $\bm{p}_{k}$ is the 3D coordinate of the $k$th input point or voxel grid, and $\bm{f}_{k}$ is the feature corresponding to $\bm{p}_{k}$. Both voxel-based and point-based convolution can then be formulated as

During the convolution, we iterate the center $\bm{x}_{k}$ over the entire input. For each center, we first index its neighbors $\bm{x}_{i}$ in $\mathcal{N}(\bm{x}_{k})$, then convolve the neighboring features $\mathcal{F}(\bm{x}_{i})$ with the kernel $\mathcal{K}(\bm{x}_{k},\bm{x}_{i})$, and finally produces the corresponding output $\bm{y}_{k}$.

Voxel-based representation is regular and has good memory locality. However, it requires very high resolution in order not to lose information. When the resolution is low, multiple points are bucketed into the same voxel grid, and these points will no longer be distinguishable. A point is kept only when it exclusively occupies one voxel grid. In Figure 2(a), we analyze the number of distinguishable points and the memory consumption (during training with batch size of 16) with different resolutions. On a single GPU (with 12 GB of memory), the largest affordable resolution is 64, which will lead to 42% of information loss (i.e., non-distinguishable points). To keep more than 90% of the information, we need to double the resolution to 128, consuming 7.2$\times$ GPU memory (82.6 GB), which is prohibitive for deployment. Although the GPU memory increases cubically with the resolution, the number of distinguishable points has a diminishing return. Therefore, the voxel-based solution is not scalable.

2 Point-Based Models: Irregular Memory Access and Dynamic Kernel Overhead

Point-based 3D modeling methods are memory efficient. The initial attempt, PointNet , is also computation efficient, but it lacks the local context modeling capability. Later research improves the expressiveness of PointNet by aggregating the neighborhood information in the point domain. However, this will lead to the irregular memory access pattern and introduce the dynamic kernel computation overhead, which becomes the efficiency bottlenecks.

Dynamic Kernel Computation.

For the 3D volumetric convolutions, the kernel $\mathcal{K}(\bm{x}_{k},\bm{x}_{i})$ can be directly indexed as the relative positions of the neighbor $\bm{x}_{i}$ are fixed for different center $\bm{x}_{k}$: e.g., each axis of the coordinate offset $\bm{p}_{i}-\bm{p}_{k}$ can only be 0, $\pm$1 for the convolution with size of 3. However, for the point-based convolution, the points are scattered over the entire 3D space irregularly; therefore, the relative positions of neighbors become unpredictable, and we will have to calculate the kernel $\mathcal{K}(\bm{x}_{k},\bm{x}_{i})$ for each neighbor $\bm{x}_{i}$ on the fly. For instance, SpiderCNN leverages the third-order Taylor expansion as a continuous approximation of the kernel $\mathcal{K}(\bm{x}_{k},\bm{x}_{i})$; PointCNN permutes the neighboring points into a canonical order with the feature transformer $\mathcal{F}(\bm{x}_{i})$. Both will introduce additional matrix multiplications. Empirically, we find that for PointCNN, the overhead of dynamic kernel computation can be more than 50% (see Figure 2(b))!

In summary, the combined overhead of irregular memory access and dynamic kernel computation ranges from 55% (for DGCNN) to 88% (for PointCNN), which indicates that most computations are wasted on dealing with the irregularity of the point-based representation.

Point-Voxel Convolution

Based on our analysis on the bottlenecks, we introduce a hardware-efficient primitive for 3D deep learning: Point-Voxel Convolution (PVConv), which combines the advantages of point-based methods (i.e., small memory footprint) and voxel-based methods (i.e., good data locality and regularity).

Our PVConv disentangles the fine-grained feature transformation and the coarse-grained neighbor aggregation so that each branch can be implemented efficiently and effectively. As illustrated in Figure 3, the upper voxel-based branch first transforms the points into low-resolution voxel grids, then it aggregates the neighboring points by the voxel-based convolutions, followed by devoxelization to convert them back to points. Either voxelization or devoxelization requires one scan over all points, making the memory cost low. The lower point-based branch extracts the features for each individual point. As it does not aggregate the neighbor’s information, it is able to afford a very high resolution.

A key component of convolution is to aggregate the neighboring information to extract local features. We choose to perform this feature aggregation in the volumetric domain due to its regularity.

The scale of different point cloud might be significantly different. We therefore normalize the coordinates $\{\bm{p}_{k}\}$ before converting the point cloud into the volumetric domain. First, we translate all points into the local coordinate system with the gravity center as origin. After that, we normalize the points into the unit sphere by dividing all coordinates by $\max\lVert\bm{p}_{k}\rVert_{2}$, and we then scale and translate the points to $$. Note that the point features$\{\bm{f}_{k}\}$remain unchanged during the normalization. We denote the normalized coordinates as$\{\hat{\bm{p}}_{k}\}$.

Voxelization.

We transform the normalized point cloud $\{(\hat{\bm{p}}_{k},\bm{f}_{k})\}$ into the voxel grids $\{\bm{V}_{u,v,w}\}$ by averaging all features $\bm{f}_{k}$ whose coordinate $\hat{\bm{p}}_{k}=(\hat{\bm{x}}_{k},\hat{\bm{y}}_{k},\hat{\bm{z}}_{k})$ falls into the voxel grid $(u,v,w)$:

Feature Aggregation.

After converting the points into voxel grids, we apply a stack of 3D volumetric convolutions to aggregate the features. Similar to conventional 3D models, we apply the batch normalization and the nonlinear activation function after each 3D convolution.

Devoxelization.

As we need to fuse the information with the point-based feature transformation branch, we then transform the voxel-based features back to the domain of point cloud. A straightforward implementation of the voxel-to-point mapping is the nearest-neighbor interpolation (i.e., assign the feature of a grid to all points that fall into the grid). However, this will make the points in the same voxel grid always share the same features. Therefore, we instead leverage the trilinear interpolation to transform the voxel grids to points to ensure that the features mapped to each point are distinct.

As our voxelization and devoxelization are both differentiable, the entire voxel-based feature aggregation branch can then be optimized in an end-to-end manner.

2 Point-Based Feature Transformation

The voxel-based feature aggregation branch fuses the neighborhood information in a coarse granularity. However, in order to model finer-grained individual point features, low-resolution voxel-based methods alone might not be enough. To this end, we directly operate on each point to extract individual point features using an MLP. Though simple, the MLP outputs distinct and discriminative features for each point. Such high-resolution individual point information is very critical to supplement the coarse-grained voxel-based information.

3 Feature Fusion

With both individual point features and aggregated neighborhood information, we can efficiently fuse two branches with an addition as they are providing complementary information.

4 Discussions

Our PVConv is more efficient than conventional point-based convolutions due to its better data locality and regularity. Our proposed voxelization and devoxelization both require $\mathcal{O}(n)$ random memory accesses, where $n$ is the number of points, since we only need to iterate over all points once to scatter them to their corresponding voxel grids. However, for conventional point-based methods, gathering the neighbors for all points requires at least $\mathcal{O}(kn)$ random memory accesses, where $k$ is the number of neighbors. Therefore, our PVCNN is $k\times$ more efficient from this viewpoint. As the typical value for $k$ is 32/64 in PointNet++ and 16 in PointCNN , we empirically reduce the number of incontiguous memory accesses by 16$\times$ to 64$\times$ through our design and achieve better data locality. Besides, as our convolutions are done in the voxel domain, which is regular, our PVConv does not require KNN computation and dynamic kernel computation, which are usually quite expensive.

Effectiveness: Keeping Points in High Resolution.

As our point-based feature extraction branch is implemented as MLP, a natural advantage is that we are able to maintain the same number of points throughout the whole network while still having the capability to model neighborhood information. Let us make a comparison between our PVConv and set abstraction (SA) module in PointNet++ . Suppose we have a batch of 2048 points with 64-channel features (with batch size of 16). We consider to aggregate information from 125 neighbors of each point and transform the aggregated feature to output the features with the same size. The SA module will require 75.2 ms of latency and 3.6 GB of memory consumption, while our PVConv will only require 25.7 ms of latency and 1.0 GB of memory consumption. The SA module will have to downsample to 685 points (i.e., around 3$\times$ downsampling) to match up with the latency of our PVConv, while the memory consumption will still be 1.5$\times$ higher. Thus, with the same latency, our PVConv is capable of modeling the full point cloud, while the SA module has to downsample the input aggressively, which will inevitably induce information loss. Therefore, our PVCNN is more effective compared to its point-based counterpart.

Experiments

We experimented on multiple 3D tasks including object part segmentation, indoor scene segmentation and 3D object detection. Our PVCNN achieves superior performance on all these tasks with lower measured latency and GPU memory consumption. More details are provided in the appendix.

We first conduct experiments on the large-scale 3D object dataset, ShapeNet Parts . For a fair comparison, we follow the same evaluation protocol as in Li et al. and Graham et al. . The evaluation metric is mean intersection-over-union (mIoU): we first calculate the part-averaged IoU for each of the 2874 test models and average the values as the final metrics. Besides, we report the measured latency and GPU memory consumption on a single GTX 1080Ti GPU to reflect the efficiency. We ensure the input data to have the same size with 2048 points and batch size of 8.

Models.

We build our PVCNN by replacing the MLP layers in PointNet with our PVConv layers. We adopt PointNet , RSNet , PointNet++ (with multi-scale grouping), DGCNN , SpiderCNN and PointCNN as our point-based baselines. We reimplement 3D-UNet as our voxel-based baseline. Note that most baselines make their implementation publicly available, and we therefore collect the statistics from their official implementation.

Results.

As in Table 1, our PVCNN outperforms all previous models. PVCNN directly improves the accuracy of its backbone (PointNet) by 2.5% with even smaller overhead compared with PointNet++. We also design narrower versions of PVCNN by reducing the number of channels to 25% (denoted as 0.25$\times$C) and 50% (denoted as 0.5$\times$C). The resulting model requires only 53.5% latency of PointNet, and it still outperforms several point-based methods with sophisticated neighborhood aggregation including RSNet, PointNet++ and DGCNN, which are almost an order of magnitude slower.

In Figure 4, PVCNN achieves a significantly better accuracy vs. latency trade-off compared with all point-based methods. With similar accuracy, our PVCNN is 15$\times$ faster than SpiderCNN and 2.7$\times$ faster than PointCNN. Our PVCNN also achieves a significantly better accuracy vs. memory trade-off compared with modern voxel-based baseline. With better accuracy, PVCNN saves the GPU memory consumption by 10$\times$ compared with 3D-UNet.

Furthermore, we also measure the latency of PVCNN on three edge devices. In Figure 5, PVCNN consistently achieves a speedup of 2$\times$ over PointNet and PointCNN on different devices. Especially, PVCNN is able to run at 19.9 objects per second on Jetson Nano with PointNet++-level accuracy and 20.2 objects per second on Jetson Xavier with PointCNN-level accuracy.

Analysis.

Conventional voxel-based methods have saturated the performance as the input resolution increases, but the memory consumption grows cubically. PVCNN is much more efficient, and the memory increases sub-linearly (Table 3). By increasing the resolution from 16 (0.5$\times$R) to 32 (1$\times$R), the GPU memory usage is increased from 1.55 GB to 1.59 GB, only 1.03$\times$. Even if we squeeze the volumetric resolution to 16 (0.5$\times$R), our method still outperforms 3D-UNet that has much higher voxel resolution (96) by a large margin (1%). PVCNN is very robust even with small resolution in the voxel branch, thanks to the high-resolution point-based branch maintaining the individual point’s information. We also compared different implementations of devoxelization in Table 3. The trilinear interpolation performs better than the nearest neighbor, which is because the points near the voxel boundaries will introduce larger fluctuations to the gradient, making it harder to optimize.

Visualization.

We illustrate the voxel and point branch features from the final PVConv in Figure 6, where warmer color represents larger magnitude. We can see that the voxel branch captures large, continuous parts (e.g. table top, lamp head) while the point branch captures isolated, discontinuous details (e.g., table legs, lamp neck). The two branches provide complementary information and can be explained by the fact that the convolution operation extracts features with continuity and locality.

2 Indoor Scene Segmentation

We conduct experiments on the large-scale indoor scene segmentation dataset, S3DIS . We follow Tchapmi et al. and Li et al. to train the models on area 1,2,3,4,6 and test them on area 5 since it is the only area that does not overlap with any other area. Both data processing and evaluation protocol are the same as PointCNN for fair comparison. We measure the latency and memory consumption with 32768 points per batch at test time on a single GTX 1080Ti GPU.

Models.

Apart from PVCNN (which is based on PointNet), we also extend PointNet++ with our PVConv to build PVCNN++. We compare our two models with the state-of-the-art point-based models and the voxel-based baseline .

Results.

As in Table 4, PVCNN improves its backbone (PointNet) by more than 13% in mIoU, and it also outperforms DGCNN (which involves sophisticated graph convolutions) by a large margin in both accuracy and latency. Remarkably, our PVCNN++ outperforms the state-of-the-art point-based model (PointCNN) by 1.7% in mIoU with 4$\times$ lower latency, and the voxel-based baseline (3D-UNet) by 4% in mIoU with more than 8$\times$ lower latency and GPU memory consumption.

Similar to object part segmentation, we design compact models by reducing the number of channels in PVCNN to 12.5%, 25% and 50% and PVCNN++ to 50%. Remarkably, the narrower version of our PVCNN outperforms DGCNN with 15$\times$ measured speedup, and RSNet with 9$\times$ measured speedup. Furthermore, it achieves 4% improvement in mIoU upon PointNet while still being 2.5$\times$ faster than this extremely efficient model (which does not have any neighborhood aggregation).

3 3D Object Detection

We finally conduct experiments on the driving-oriented dataset, KITTI . We follow Qi et al. to construct the val set from the training set so that no instances in the val set belong to the same video clip of any training instance. The size of val set is 3769, leaving the other 3711 samples for training. We evaluate all models for 20 times and report the mean 3D average precision (AP).

Models.

We build two versions of PVCNN based on F-PointNet : (a) an efficient version where we only replace the MLP layers within the instance segmentation network, and (b) a complete version where we further replace the MLP layers in the box estimation network. We compare our two models with F-PointNet (whose backbone is PointNet) and F-PointNet++ (whose backbone is PointNet++).

Results.

In Table 5, even if our efficient model does not aggregate neighboring features in the box estimation network while F-PointNet++ does, ours still outperform it in most classes with 1.8$\times$ lower latency. Improving the box estimation network with PVConv, our complete model outperforms both baselines in all categories significantly. Compared with F-PointNet baseline, our PVCNN obtains up to 8% mAP improvement in pedestrians and 3.5-6.8% mAP improvement in cyclist, which indicates that our proposed PVCNN is both efficient and expressive.

Conclusion

We propose Point-Voxel CNN (PVCNN) for fast and efficient 3D deep learning. We bring the best of both worlds together: voxels and points, reducing the memory footprint and irregular memory access. We represent the 3D input data efficiently with the sparse, irregular point representation and perform the convolutions efficiently in the dense, regular voxel representation. Extensive experiments on multiple tasks consistently demonstrate the effectiveness and efficiency of our proposed method. We believe that our research will break the stereotype that the voxel-based convolution is naturally inefficient and shed light on co-designing the voxel-based and point-based architectures for fast and efficient 3D deep learning.

We thank MIT Quest for Intelligence, MIT-IBM Watson AI Lab, Samsung, Facebook and SONY for supporting this research. We thank AWS Machine Learning Research Awards for providing the computation resource. We thank NVIDIA for donating Jetson AGX Xavier.

References