Axial Attention Inside a U-Net for Semantic Segmentation of 3D Sparse LiDAR Point Clouds

Semantic segmentation of LiDAR point clouds is essential for driving guidance in autonomous driving. The current approaches may use complex data representations for sparse LiDAR point clouds or sophisticated neural networks to process these data; however, such methods may not be suitable for real-time prediction on mobile devices due to the high computational requirements.In this research, we propose to enhance U-Nets with axial attention modules to process the point clouds on 2D planes after the spherical projection. U-Nets are simple and baseline encoder-decoder networks for semantic segmentation with convolution operations of local context. We improved baseline U-Nets by simply adding axial attention modules inside U-Nets to increase the feature extraction effectiveness by incorporating the capability of modeling long-range dependencies from axial attention.Experiments were performed on the benchmark KITTI and our lab LiDAR datasets. In terms of mIoU (mean intersection over union), performance was increased from 51.3% of U-Nets to 54.8% of the mixed axial-attention U-Nets in KITTI, and from 81.7% of U-Nets to 92.3% of the mixed axial-attention U-Nets in our LiDAR dataset. In testing, the latency of the mixed axial-attention U-Nets in our LiDAR dataset was 22.3 ms per scan, which was fast enough to be real-time.We propose to embed axial attention in the decoder of U-Nets for semantic segmentation of 3D LiDAR point clouds with better IoU performance and real-time prediction capability on mobile devices.