PointDistiller: Structured Knowledge Distillation Towards Efficient and Compact 3D Detection
Linfeng Zhang, Runpei Dong, Hung-Shuo Tai, Kaisheng Ma
Introduction
The growth in large-scale lidar datasets and the achievements in end-to-end 3D representation learning have boosted the developments of point cloud based segmentation, generation, and detection . As one of the essential tasks of 3D computer vision, 3D object detection plays a fundamental role in real-world applications such as autonomous driving cars and virtual reality . However, recent research has shown a growing discrepancy between cumbersome 3D detectors that achieve state-of-the-art performance and lightweight 3D detectors which are affordable in real-time applications on edge devices. To address this problem, sufficient model compression techniques have been proposed, such as network pruning , quantization , lightweight model design , and knowledge distillation .
Knowledge distillation, which aims to improve the performance of a lightweight student model by training it to mimic a pre-trained and over-parameterized teacher model, has evolved into one of the most popular and effective model compression methods in both computer vision and natural language processing . Sufficient theoretical and empirical results have demonstrated its effectiveness in image-based visual tasks such as image classification , semantic segmentation and object detection . However, compared with images, point clouds have their properties: (i) Point clouds inherently lack topological information, which makes recovering the local topology information crucial for the visual tasks . (ii) Different from images that have a regular structure, point clouds are irregularly and sparsely distributed in the metric space .
These differences between images and point clouds have hindered the image-based knowledge distillation methods from achieving satisfactory performance on point clouds and also raised the requirement to design specific knowledge distillation methods for point clouds. Recently, a few methods have been proposed to apply knowledge distillation to 3D detection . However, most of these methods focus on the choice of student-teacher in a multi-modal setting, e.g., teaching point clouds-based student detectors with an images-based teacher or vice versa, and still ignore the peculiar properties of point clouds. To address this problem, we propose a structured knowledge distillation framework named PointDistiller, which involves local distillation to distill teacher knowledge in the local geometric structure of point clouds, and reweighted learning strategy to handle the sparsity of point clouds by highlighting student learning on the relatively more crucial voxels.
Local Distillation Sufficient recent studies show that capturing and making usage of the semantic information in the local geometric structure of point clouds have a crucial impact on point cloud representation learning . Hence, instead of directly distilling the backbone feature of teacher detectors to student detectors, we propose local distillation, which firstly clusters the local neighboring voxels or points with KNN (K-Nearest Neighbours), then encodes the semantic information in local geometric structure with dynamic graph convolutional layers , and finally distill them from teachers to students. Hence, the student detectors can inherit the teacher’s ability to understand point clouds’ local geometric information and achieve better detection performance.
Reweighted Learning Strategy One of the mainstream methods for processing point clouds is to convert them into volumetric voxels and then encode them as regular data. However, due to the sparsity and the noise in point clouds, most of these voxels contain only a single point. For instance, as shown in Figure 2, on the KITTI dataset, around 68% voxels in point clouds contain only one point, which has a high probability of being a noise point. Hence, the representative features in these single-point voxels have relatively lower importance in knowledge distillation compared with the voxels which contain multiple points. Motivated by this observation, we propose a reweighted learning strategy, which highlights student learning on the voxels with multiple points by giving them larger learning weights. Besides, the similar idea can also be easily extended to raw points-based detectors to highlight knowledge distillation on the points which have a more considerable influence on the prediction of the teacher detector.
Extensive experiments on both voxels-based and raw-points based detectors have been conducted to demonstrate the effectiveness of our method over the previous seven knowledge distillation methods. As shown in Figure 2, on PointPillars and SECOND detectors, our method leads to 4 compression and 0.91.8 mAP improvements at the same time. On PointRCNN, our method leads to 8 compression with only 0.2 BEV mAP drop. Our main contributions be summarized as follows.
We propose local distillation, which firstly encodes the local geometric structure of point clouds with dynamic graph convolution and then distills them from teachers to students.
We propose reweighted learning strategy to handle the sparsity and noise in point clouds. It highlights student learning on the voxels, which have more points inside them, by giving them higher learning weights in knowledge distillation.
Extensive experiments on both voxels-based and raw points-based detectors have been conducted to demonstrate the performance of our method over seven previous methods. Besides, we have released our codes to promote future research.
Related Work
The idea of training a small model with a large pre-trained model is firstly proposed by Buciluǎ et al. for ensemble model compression . Then, with the excellent breakthroughs of deep learning, Hinton et al. propose the concept of knowledge distillation which strives to compress an over-parameterized teacher model by transferring its knowledge to a lightweight student model . Early knowledge distillation methods usually train the students to mimic the predicted categorical probability distribution of teachers . Then, extensive methods have been proposed to learn teacher knowledge in the backbone features or its variants, such as attention , relation , task-oriented information and so on. Following its success in classification, abundant works have applied knowledge distillation to object detection , segmentation , image generation , pre-trained language models , semi-supervised learning and lead to consistent effectiveness.
Recently, designing specific knowledge distillation methods to improve the efficiency and accuracy of object detection has become a rising and popular topic. Chen et al. first propose to apply the naive prediction and feature-based knowledge distillation methods to object detection . Then, Wang et al. show that the imbalance between foreground objects and background objects hinders knowledge distillation from achieving better performance in object detection . To address this problem, abundant knowledge distillation methods have tried to find the to-be-distilled regions based on the ground-truth , detection results , spatial attention , query-based attention and gradients . Moreover, recent methods have also been proposed to distill the pixel-level and object-level relation from teachers to students . Besides knowledge distillation for 2D detection, some cross-modal knowledge distillation have been introduced to transfer knowledge from RGB-based teacher detectors to lidar-based student detectors or vice versa . However, most of these methods focus on the choice of students and teachers in a multi-modal framework, while the design of specific knowledge distillation optimization methods on point clouds based pure 3D detection has not been well-explored.
2 3D Object Detection on Point Clouds
The rapid development of deep learning has firstly boosted the research in 2D object detection and then recently raised the research trend in point clouds-based 3D object detection. PointNet is firstly proposed to extract the feature of points with multi-layer perception in an end-to-end manner. Then, PointNet++ is further proposed to capture the local structures in a hierarchical fashion with density adaptive sampling and grouping . Zhou et al. propose VoxelNet, a single-stage detector that divides a point cloud into equally spaced 3D voxels and processes them with voxel feature encoding layers . Then, SECOND is proposed to improve VoxelNet with sparse convolutional layers and focal loss . PointPillars is proposed to divide point clouds into several pillars and then convert them into a pseudo image, which can be further processed with 2D convolutional layers . Shi et al. propose PointRCNN, a two-stage detection method that firstly generates bottom-up 3D proposals based on the raw point clouds and then refines them to obtain the final detection results . Afterward, Fast Point R-CNN and PV-RCNN are proposed to utilize both voxel representation and raw point clouds to exploit their respective advantages . Recently, Qi et al. propose to perform offboard 3D detection with point cloud sequences, which is able to make use of the temporal points and achieve comparable performance with human labels . The graph convolutional neural network is another rising star in point cloud detection . Lin et al. propose 3D-GCN to avoid the shift and scale changes in point clouds . Zhou et al. propose adaptive graph convolution, which generates adaptive kernels according to the learned features .
Unfortunately, the significant 3D detection performance usually comes at the expense of high computational and storage costs, making them unaffordable in real-time applications such as self-driving cars. To address this issue, recent research attention has been paid to designing efficient 3D detectors. Tang et al. propose to apply neural architecture search to 3D detection by using sparse point-voxel convolution . Li et al. propose Lidar-RCNN, which resorts to a point-based approach and remedies the problem of uncorrected proposal sizes . Liu et al. propose voxel-point cnn to represent the 3D input data in points while performing the convolutions in voxels to reduce the memory accessing consumption . Recently, Li et al. propose to improve the efficiency of graph convolution for point clouds by simplified KNN search and graph shuffling .
Methodology
2 Our Method
Extracting Local Geometric Information
As pointed out by abundant previous works, the local geometric information has a crucial influence on the performance of point cloud detectors . Thus, instead of directly distilling the representative feature, we propose local distillation which extracts the local geometric information of point clouds with dynamic graph convolution layers and distills it to the student detector. Concretely, denoting as the feature of the -th to-be-distilled voxel or point, we can build a graph based on this voxel or point and its neighboring voxels or points clustered by KNN (K-Nearest Neighbours). By denoting the features of and its neighbours as and respectively, motivated by previous methods , we firstly update the feature of each voxel (or point) in this graph by concatenating them with the global centroid voxel (or point) feature , which can be formulated as \hat{z}_{i,j}=\text{cat}\big{(}[z_{i,1},z_{i,j}]\big{)} for all . Then, we apply a dynamic graph convolution as the aggregation operation upon them, which can be formulated as , where is the aggregation operator. Following previous graph-based point cloud networks, we set as a nonlinear layer with ReLU activation and batch normalization. Then the training objective of local distillation can be formulated as
where indicates the parameters of student encoding layer . indicates the parameters of dynamic graph convolution layers for the student and teacher detectors. Note that these layers are trained with the student detector simultaneously and can be discarded during inference.
Reweighting Knowledge Distillation Loss
As shown in the above loss function, with a higher , the knowledge distillation loss between student and teacher features at the -th graph will have a more extensive influence on the overall loss, and thus student learning on the -th graph can be highlighted. As a result, the proposed reweighting strategy allows the student detector to pay more attention to learning teacher knowledge in the relatively more crucial voxel graphs (point graphs). Moreover, Equation 2 also implies that our method is a feature-based knowledge distillation method that is not correlated with the architecture of detectors and the label set . Hence, it can be directly added to the origin training loss of all kinds of 3D object detectors for model compression.
Experiment
We have evaluated our method in both voxels-based object detector including PointPillars and SECOND , and the raw points based object detector including PointRCNN . Most experiments are conducted on KITTI and nuScenes , which consist of samples that have both lidar point clouds and images. Our models are trained with only the lidar point clouds. For KITTI, we report the average precision calculated by 40 sampling recall positions for BEV (Bird’s Eye View) object detection and 3D object detection on the validation split. Following the typical protocol, the IoU threshold is set as 0.7 for class Car and 0.5 for class Pedestrians and Cyclists. We have mainly compared our methods with seven previous knowledge distillation methods, including methods proposed by Remero et al. , Zagoruko et al. , Tung et al. , Heo et al. , Zheng et al. , Tian et al. , and Zhang et al. . All the experiments are conducted with mmdetection3d and PyTorch . We keep the training and evaluation settings in mmdetection3d as default. The teacher model is the origin model before compression. The student model shares the same architecture and depth as its teacher but with fewer channels. Following previous works, the average precision of three difficulties and the three categories are reported as the performance metrics . Please refer to our codes in the supplementary material for more details.
2 Experimental Results
Table 1 and Table 2 show the performance of detectors trained with and without our method for BEV detection and 3D detection, respectively. It is observed that: (i) Significant average precision improvements on all kinds of detectors and all compression ratios for both BEV and 3D detection. On average, 2.4 and 1.0 moderate mAP improvements can be observed for the voxel and raw points-based detectors, respectively. On BEV and 3D detection, 1.9 and 1.9 moderate mAP improvements can be obtained, respectively. (ii) On the BEV detection of PointPillars and SECOND detectors, the 4 compressed and accelerated students trained with our method outperform their teachers by 0.9 and 0.9 mAP, respectively. On the 3D detection of PointPillars and SECOND detectors, the 4 compressed and accelerated students trained with our method outperform their teachers by 1.8 and 0.1 mAP, respectively. (iii) Consistent average precision boosts can be observed in detection results of all difficulties. For instance, on BEV detection of PointPillars students, 2.4, 2.3, and 2.3 mAP improvements can be observed for easy, moderate, and hard difficulties, respectively. These observations demonstrate that our method can successfully transfer teacher knowledge to the student detectors. (iv) Consistent average precision boosts can be observed in detection results of all categories. For instance, on moderate BEV detection of PointPillars students, 0.6, 3.2 and 3.1 mAP improvements can be obtained on cars, pedestrians and cyclists, respectively. (v) On PointRCNN, on average 1.3 and 1.2 moderate mAP improvements can be observed on BEV and 3D detection, respectively, indicating that our method is also effective for raw points-based detectors. In summary, these experiment results demonstrate that our method can successfully transfer the knowledge from teacher detectors to student detectors and lead to significant and consistent performance boosts.
Comparison between our method and previous knowledge distillation methods is shown in Table 3. It is observed that: (i) Our method outperforms the previous methods by a clear margin. On BEV and 3D detection, our method outperforms the second-best knowledge distillation method by 1.5 and 1.9 moderate mAP, respectively. (ii) Our method achieves the best performance for all categories of all difficulties. (iii) Besides, our method is the only knowledge distillation method that enables the student detector to outperform its teacher detector.
Experiments on nuScenes
Experiments of 2 and 4 compressed PointPillars on nuScenes are shown in Table 4. It is observed that our method leads to 0.65 and 0.5 improvements on mAP and NDS on average, respectively, indicating that our method is also effective on the large-scale dataset.
Discussion
The proposed PointDistiller is mainly composed of two components, including the reweighted learning strategy (RL) and local distillation (LD). Ablation studies with 4 compressed PointPillars students on KITTI are shown in Table 5. It is observed that: (i) 2.0 and 1.9 mAP improvements can be obtained by only using the reweighted learning strategy to distill the backbone features on BEV detection and 3D detection, respectively. (ii) 2.3 and 2.5 mAP boosts can be gained by using local distillation without reweighted learning on BEV detection and 3D detection, respectively. (iii) By combining the two methods together, 0.5 and 0.9 further mAP improvements can be achieved on BEV detection and 3D detection, respectively. These observations indicate that each module in PointDistiller has its individual effectiveness and their merits are orthogonal. Besides, they also implies that the proposed local distillation and reweighted learning may be combined with other knowledge distillation methods to achieve better performance.
Sensitivity Study
Our method mainly introduces two hyper-parameters, , and , which indicate the number of nodes in a graph for local distillation, and the number of to-be-distilled voxels (points) respectively. A hyper-parameter sensitivity study on the two hyper-parameters is shown in Figure 6. It is observed that our method with different hyper-parameter values consistently outperforms the baseline by a large margin, indicating our method is not sensitive to hyper-parameters.
2 Visualization Analysis
In the reweighted learning strategy, the importance scores of each voxel or point are utilized to determine whether it should be distilled. Visualization of the importance scores in PointPillars is shown in Figure 6. It is observed that they successfully localize the foreground objects (e.g., cars and pedestrians) and the hard-negative objects (e.g., walls).
Visualization on Detection Results
In this subsection, we have visualized the detection results of the student model trained with and without our method for comparison. Note that both student models are 4 compressed PointPillars trained on KITTI. The green and blue boxes indicate the boxes of the model prediction and the ground truth. As shown in Figure 7, the student model without knowledge distillation tends to have much more false-positive (FP) predictions. In contrast, this excessive FP problem is alleviated in the student trained with our method. This observation is consistent with our experimental results that the distilled PointPillars has 3.4 mAP improvements.
Conclusion
This paper proposes a structured knowledge distillation framework named PointDistiller for point clouds-based object detection. It is composed of local distillation to first encode the semantic information in local geometric structure in point clouds and distill it to students, and reweighted learning to handle the sparsity and noise in point clouds by assigning different learning weights to different points and voxels. Extensive experiments on both voxels-based detectors and raw points-based detectors have demonstrated the superiority over seven previous knowledge distillation methods. Our ablation study has shown the individual effectiveness of each module in PointDistiller. Besides, the visualization results demonstrate that PointDistiller can significantly improve detection performance by reducing false-positive predictions, and the importance score is able to reveal the more significant voxels. To the best of our knowledge, this work initiates the first step to exploring KD for efficient point clouds-based 3D object detection, and we hope this could spur future research.