FlatFormer：高效点云Transformer的扁平化窗口注意力

网站 | 论文

摘要

Transformer作为CNN的替代方案，已在多种模态（如文本和图像）中证明其有效性。对于3D点云Transformer，现有研究主要集中于将其准确性提升到最先进水平。然而，它们的延迟落后于基于稀疏卷积的模型（慢3倍），阻碍了它们在资源受限、对延迟敏感的应用（如自动驾驶）中的使用。这种低效源于点云的稀疏和不规则性质，而Transformer设计用于处理密集、规则的工作负载。本文提出FlatFormer，通过牺牲空间邻近性来换取更好的计算规则性，从而缩小这一延迟差距。我们首先使用基于窗口的排序来扁平化点云，并将点分成大小相等的组而不是形状相等的窗口。这有效避免了昂贵的结构化和填充开销。然后，我们在组内应用自注意力提取局部特征，交替排序轴以收集不同方向的特征，并移动窗口以在组间交换特征。FlatFormer在Waymo开放数据集上实现了最先进的准确性，比（基于Transformer的）SST快4.6倍，比（稀疏卷积的）CenterPoint快1.4倍。这是首个在边缘GPU上实现实时性能的点云Transformer，比稀疏卷积方法更快，同时在大规模基准测试上达到相当甚至更优的准确性。

结果

本仓库中的所有结果均可复现。遗憾的是，由于Waymo数据集许可协议，我们无法提供预训练模型权重。如果您在项目中使用FlatFormer无法获得令人满意的性能，欢迎讨论。

3D目标检测（在Waymo验证集上）

使用方法

先决条件

代码基于以下库构建：

• Python >= 3.6, <3.8
• PyTorch >= 1.9, <= 1.10.2
• mmcv = 1.4.0
• mmdetection = 2.14.0

安装这些依赖项后，请运行以下命令安装代码库：

pip install -v -e .

数据集准备

请按照MMDetection3D的说明下载并预处理Waymo开放数据集。完成数据准备后，您将看到以下目录结构（如MMDetection3D中所示）：

mmdetection3d
├── mmdet3d
├── tools
├── configs
├── data
│   ├── waymo
│   │   ├── waymo_format
│   │   │   ├── training
│   │   │   ├── validation
│   │   │   ├── testing
│   │   │   ├── gt.bin
│   │   ├── kitti_format
│   │   │   ├── ImageSets
│   │   │   ├── training
│   │   │   ├── testing
│   │   │   ├── waymo_gt_database
│   │   │   ├── waymo_infos_trainval.pkl
│   │   │   ├── waymo_infos_train.pkl
│   │   │   ├── waymo_infos_val.pkl
│   │   │   ├── waymo_infos_test.pkl
│   │   │   ├── waymo_dbinfos_train.pkl

训练

# 多GPU训练
bash tools/dist_train.sh configs/flatformer/$CONFIG.py 8 --work-dir $CONFIG/ --cfg-options evaluation.pklfile_prefix=./work_dirs/$CONFIG/results evaluation.metric=waymo

评估

# 多GPU测试
bash tools/dist_test.sh configs/flatformer/$CONFIG.py /work_dirs/$CONFIG/latest.pth 8 --eval waymo

引用

如果FlatFormer对您的研究有用或相关，请通过引用我们的论文来认可我们的贡献：

@inproceedings{liu2023flatformer,
  title={FlatFormer: Flattened Window Attention for Efficient Point Cloud Transformer},
  author={Liu, Zhijian and Yang, Xinyu and Tang, Haotian and Yang, Shang and Han, Song},
  booktitle={IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2023}
}

致谢

本项目基于以下代码库：

• MMDetection3D
• SST

我们要感谢Tianwei Yin、Lue Fan和Ligeng Mao提供CenterPoint、SST/FSD和VoTr的详细结果，以及Yue Wang和Yukang Chen的有益讨论。本工作得到了国家科学基金会、MIT-IBM Watson人工智能实验室、NVIDIA、现代和福特的支持。Zhijian Liu部分受到高通创新奖学金的资助。