超点变换器
官方实现
使用超点变换器进行高效的3D语义分割 (ICCV 2023)
可扩展的3D全景分割作为超点图聚类 (3DV 2024 口头报告)
如果你❤️或仅仅使用了这个项目,别忘了给这个仓库一个⭐,
这对我们意义重大!
@article{robert2023spt,
title={Efficient 3D Semantic Segmentation with Superpoint Transformer},
author={Robert, Damien and Raguet, Hugo and Landrieu, Loic},
journal={Proceedings of the IEEE/CVF International Conference on Computer Vision},
year={2023}
}
@article{robert2024scalable,
title={Scalable 3D Panoptic Segmentation as Superpoint Graph Clustering},
author={Robert, Damien and Raguet, Hugo and Landrieu, Loic},
journal={Proceedings of the IEEE International Conference on 3D Vision},
year={2024}
}
📌 描述
超点变换器
超点变换器 (SPT) 是一种基于超点的变换器🤖架构,可以高效地⚡对大规模3D场景进行语义分割。该方法包括一种快速算法,将点云划分🧩为分层的超点结构,以及一种自注意力机制,用于利用多个尺度上超点之间的关系。
✨ SPT数据 ✨ |
---|
📊 S3DIS 6折交叉验证上的最佳性能 (76.0 mIoU) |
📊 KITTI-360验证集上的最佳性能 (63.5 mIoU) |
📊 DALES上接近最佳性能 (79.6 mIoU) |
🦋 212k参数 (PointNeXt ÷ 200, Stratified Transformer ÷ 40) |
⚡ S3DIS训练在1个GPU上3小时内完成 (PointNeXt ÷ 7, Stratified Transformer ÷ 70) |
⚡ 预处理速度比SPG快7倍 |
SuperCluster
SuperCluster是一种基于SPT的超点架构,用于(非常)大型3D场景🐘的全景分割。 我们将全景分割任务表述为可扩展的超点图聚类任务。 为此,我们的模型被训练以预测图优化问题的输入参数,其解决方案是全景分割💡。 这种表述允许仅使用每个节点和每条边的目标来监督我们的模型,避免了在训练时计算实际全景分割和相关匹配问题的需求。 在推理时,我们快速并行化的算法解决了小规模图优化问题,产生对象实例👥。 由于其轻量级骨干网络和可扩展的表述,SuperCluster可以一次性在单个GPU🚀上处理前所未有规模的场景,参数少于1M🦋。
✨ SuperCluster数据 ✨ |
---|
📊 S3DIS 6折交叉验证上的最佳性能 (55.9 PQ) |
📊 S3DIS Area 5上的最佳性能 (50.1 PQ) |
📊 ScanNet验证集上的最佳性能 (58.7 PQ) |
📊 KITTI-360验证集上的首次最佳性能 (48.3 PQ) |
📊 DALES上的首次最佳性能 (61.2 PQ) |
🦋 212k参数 (PointGroup ÷ 37) |
⚡ S3DIS训练在1个GPU上4小时内完成 |
⚡ 7.8km² 大小、1800万个点的图块在1个GPU上10.1秒内处理完成 |
📰 更新
- 2024年6月27日 发布了我们的超点变换器🧑🏫教程 幻灯片、 笔记本和视频。 如果你刚开始接触这个项目,一定要看看这些!
- 2024年6月21日 Damien将在📅 2024年6月27日中欧时间下午1点 进行一场🧑🏫超点变换器教程。 如果你想获得这个项目的实践经验,请务必参加! 在此注册。
- 2024年6月14日 为我们的**交互式可视化工具**添加了一些文档。
- 2024年2月28日 发布了全景分割的主要代码,实现了
可扩展的3D全景分割作为超点图聚类。
这个新版本还实现了期待已久的功能,如lightning的
predict()
行为、体素分辨率和全分辨率预测。 依赖项和仓库结构的一些变化不向后兼容。如果你已经在使用之前的代码版本,我们建议重新安装你的conda环境并重新运行数据集的预处理❗ - 2023年10月15日 我们的论文**可扩展的3D全景分割作为超点图聚类被接受为3DV 2024的口头报告**🥳
- 2023年10月6日 来看看我们在**ICCV 2023上的使用超点变换器进行高效的3D语义分割**海报
- 2023年7月14日 我们的论文**使用超点变换器进行高效的3D语义分割被[ICCV 2023](https://iccv2023
🏗 安装
只需运行 install.sh
即可在名为 spt
的新 conda 环境中安装所有依赖项。
# 创建名为 'spt' 的 conda 环境并安装依赖项
./install.sh
注意: 有关设置数据集路径和文件结构,请参阅 数据集页面。
🔩 项目结构
└── superpoint_transformer
│
├── configs # Hydra 配置
│ ├── callbacks # 回调配置
│ ├── data # 数据配置
│ ├── debug # 调试配置
│ ├── experiment # 实验配置
│ ├── extras # 额外工具配置
│ ├── hparams_search # 超参数搜索配置
│ ├── hydra # Hydra 配置
│ ├── local # 本地配置
│ ├── logger # 日志记录器配置
│ ├── model # 模型配置
│ ├── paths # 项目路径配置
│ ├── trainer # 训练器配置
│ │
│ ├── eval.yaml # 评估主配置
│ └── train.yaml # 训练主配置
│
├── data # 项目数据(参见 docs/datasets.md)
│
├── docs # 文档
│
├── logs # 由 hydra 和 lightning 日志记录器生成的日志
│
├── media # 项目说明媒体
│
├── notebooks # Jupyter 笔记本
│
├── scripts # Shell 脚本
│
├── src # 源代码
│ ├── data # 分层分区的数据结构
│ ├── datamodules # Lightning DataModules
│ ├── datasets # 数据集
│ ├── dependencies # 已编译的依赖项
│ ├── loader # 数据加载器
│ ├── loss # 损失函数
│ ├── metrics # 评估指标
│ ├── models # 模型架构
│ ├── nn # 模型构建块
│ ├── optim # 优化
│ ├── transforms # 转换、预处理等函数
│ ├── utils # 实用工具
│ ├── visualization # 交互式可视化工具
│ │
│ ├── eval.py # 运行评估
│ └── train.py # 运行训练
│
├── tests # 各种测试
│
├── .env.example # 存储私有环境变量的示例文件
├── .gitignore # git 忽略的文件列表
├── .pre-commit-config.yaml # 代码格式化的预提交钩子配置
├── install.sh # 安装脚本
├── LICENSE # 项目许可证
└── README.md
注意: 有关
data/
的更多详细信息,请参阅 数据集页面。
注意: 有关
logs/
的更多详细信息,请参阅 日志页面。
🚀 使用方法
数据集
请参阅 数据集页面 来设置您的数据集。
评估
使用以下命令结构从检查点文件 checkpoint.ckpt
评估我们的模型,其中 <task>
应为使用 SPT 的 semantic
或使用 SuperCluster 的 panoptic
:
# 在 <dataset> 上评估 <task> 分割
python src/eval.py experiment=<task>/<dataset> ckpt_path=/path/to/your/checkpoint.ckpt
一些示例:
# 在 S3DIS Fold 5 上评估 SPT
python src/eval.py experiment=semantic/s3dis datamodule.fold=5 ckpt_path=/path/to/your/checkpoint.ckpt
# 在 KITTI-360 Val 上评估 SPT
python src/eval.py experiment=semantic/kitti360 ckpt_path=/path/to/your/checkpoint.ckpt
# 在 DALES 上评估 SPT
python src/eval.py experiment=semantic/dales ckpt_path=/path/to/your/checkpoint.ckpt
# 在 S3DIS Fold 5 上评估 SuperCluster
python src/eval.py experiment=panoptic/s3dis datamodule.fold=5 ckpt_path=/path/to/your/checkpoint.ckpt
# 在 S3DIS Fold 5 上评估 SuperCluster,将 {wall, floor, ceiling} 作为 'stuff'
python src/eval.py experiment=panoptic/s3dis_with_stuff datamodule.fold=5 ckpt_path=/path/to/your/checkpoint.ckpt
# 在 ScanNet Val 上评估 SuperCluster
python src/eval.py experiment=panoptic/scannet ckpt_path=/path/to/your/checkpoint.ckpt
# 在 KITTI-360 Val 上评估 SuperCluster
python src/eval.py experiment=panoptic/kitti360 ckpt_path=/path/to/your/checkpoint.ckpt
# 在 DALES 上评估 SuperCluster
python src/eval.py experiment=panoptic/dales ckpt_path=/path/to/your/checkpoint.ckpt
注意:
SPT 和 SPT-nano 模型在 S3DIS 6-Fold、KITTI-360 Val 和 DALES 上的预训练权重可在以下位置获取:
SuperCluster 模型在 S3DIS 6-Fold、S3DIS 6-Fold with stuff、ScanNet Val、KITTI-360 Val 和 DALES 上的预训练权重可在以下位置获取:
训练
使用以下命令结构在 32G-GPU 上训练我们的模型,其中 <task>
应为使用 SPT 的 semantic
或使用 SuperCluster 的 panoptic
:
# 在 <dataset> 上训练 <task> 分割
python src/train.py experiment=<task>/<dataset>
一些示例:
# 在 S3DIS Fold 5 上训练 SPT
python src/train.py experiment=semantic/s3dis datamodule.fold=5
# 在 KITTI-360 Val 上训练 SPT
python src/train.py experiment=semantic/kitti360
# 在 DALES 上训练 SPT
python src/train.py experiment=semantic/dales
# 在 S3DIS Fold 5 上训练 SuperCluster
python src/train.py experiment=panoptic/s3dis datamodule.fold=5
# 在 S3DIS Fold 5 上训练 SuperCluster,将 {wall, floor, ceiling} 作为 'stuff'
python src/train.py experiment=panoptic/s3dis_with_stuff datamodule.fold=5
# 在 ScanNet Val 上训练 SuperCluster
python src/train.py experiment=panoptic/scannet
# 在 KITTI-360 Val 上训练 SuperCluster
python src/train.py experiment=panoptic/kitti360
# 在 DALES 上训练 SuperCluster
python src/train.py experiment=panoptic/dales
使用以下命令在 11G-GPU 上训练 💾(训练时间和性能可能会有所不同):
# 在 S3DIS Fold 5 上训练 SPT
python src/train.py experiment=semantic/s3dis_11g datamodule.fold=5
# 在 KITTI-360 Val 上训练 SPT
python src/train.py experiment=semantic/kitti360_11g
# 在 DALES 上训练 SPT
python src/train.py experiment=semantic/dales_11g
# 在 S3DIS Fold 5 上训练 SuperCluster
python src/train.py experiment=panoptic/s3dis_11g datamodule.fold=5
# 在 S3DIS Fold 5 上训练 SuperCluster,将 {wall, floor, ceiling} 作为 'stuff'
python src/train.py experiment=panoptic/s3dis_with_stuff_11g datamodule.fold=5
# 在 ScanNet Val 上训练 SuperCluster
python src/train.py experiment=panoptic/scannet_11g
# 在 KITTI-360 Val 上训练 SuperCluster
python src/train.py experiment=panoptic/kitti360_11g
# 在 DALES 上训练 SuperCluster
python src/train.py experiment=panoptic/dales_11g
注意:遇到 CUDA 内存不足错误 💀💾?请查看我们专门的 故障排除部分。
注意:在
configs/experiment/
中提供了其他即用型配置。您可以通过组合 configs 轻松设计自己的实验:# 在 DALES 上训练 Nano-3 模型 50 个 epoch python src/train.py datamodule=dales model=nano-3 trainer.max_epochs=50
有关配置系统如何工作以及 Lightning+Hydra 组合的所有优秀特性的更多信息,请参阅 Lightning-Hydra。
注意:默认情况下,您的日志将自动上传到 Weights and Biases,您可以从那里跟踪和比较您的实验。其他日志记录器可在
configs/logger/
中找到。有关日志记录选项的更多信息,请参阅 Lightning-Hydra。
PyTorch Lightning predict()
SPT 和 SuperCluster 都继承自 LightningModule
并实现了 predict_step()
,这允许使用 PyTorch Lightning 的 Trainer.predict()
机制。
from src.models.semantic import SemanticSegmentationModule
from src.datamodules.s3dis import S3DISDataModule
from pytorch_lightning import Trainer
# 在 S3DIS 上进行语义分割预测
dataloader = S3DISDataModule(...)
model = SemanticSegmentationModule(...)
trainer = Trainer(...)
batch, output = trainer.predict(model=model, dataloaders=dataloader)
然而,这仍然需要您实例化一个 Trainer
、一个 DataLoader
和一个带有相关参数的模型。
为了简化一些,我们所有的数据集都继承自 LightningDataModule
并默认通过指向其相应的测试集来实现 predict_dataloader()
。这允许直接将 datamodule 传递给 PyTorch Lightning 的 Trainer.predict()
,而无需显式实例化 DataLoader
。
from src.models.semantic import SemanticSegmentationModule
from src.datamodules.s3dis import S3DISDataModule
from pytorch_lightning import Trainer
# 在 S3DIS 上进行语义分割预测
datamodule = S3DISDataModule(...)
model = SemanticSegmentationModule(...)
trainer = Trainer(...)
batch, output = trainer.predict(model=model, datamodule=datamodule)
有关如何实例化这些以及我们模型的输出格式的更多详细信息,我们强烈建议您尝试我们的 演示笔记本 并查看 src/eval.py
脚本。
全分辨率预测
根据设计,我们的模型在训练期间只需要为 $P_1$ 分区级别的超点生成预测。所有损失和指标都被制定为超点级别的目标。这在训练和评估时方便地节省了计算和内存。
然而,在推理时,我们通常需要 $P_0$ 分区级别的体素预测或全分辨率输入点云的预测。为此,我们提供了辅助函数来恢复体素级和全分辨率预测。
有关这些的更多详细信息,请参阅我们的 演示笔记本。
在自定义数据上使用预训练模型
要在自己的点云上运行预训练模型,请参阅我们的教程 幻灯片、笔记本 和 视频。
在自定义数据上参数化超点分区
我们的分层超
👩🔧 故障排除
以下是一些常见问题及解决建议。
在11G-GPU上运行SPT或SuperCluster
我们的默认配置是为32G-GPU设计的。但是,SPT和SuperCluster可以在11G-GPU 💾上运行,只是时间和性能会有轻微变化。
我们在configs/experiment/semantic
中提供了在11G-GPU 💾上训练SPT的配置:
# 在S3DIS Fold 5上训练SPT
python src/train.py experiment=semantic/s3dis_11g datamodule.fold=5
# 在KITTI-360 Val上训练SPT
python src/train.py experiment=semantic/kitti360_11g
# 在DALES上训练SPT
python src/train.py experiment=semantic/dales_11g
同样,我们在configs/experiment/panoptic
中提供了在11G-GPU 💾上训练SuperCluster的配置:
# 在S3DIS Fold 5上训练SuperCluster
python src/train.py experiment=panoptic/s3dis_11g datamodule.fold=5
# 在S3DIS Fold 5上训练SuperCluster,将{墙、地板、天花板}作为'stuff'
python src/train.py experiment=panoptic/s3dis_with_stuff_11g datamodule.fold=5
# 在ScanNet Val上训练SuperCluster
python src/train.py experiment=panoptic/scannet_11g
# 在KITTI-360 Val上训练SuperCluster
python src/train.py experiment=panoptic/kitti360_11g
# 在DALES上训练SuperCluster
python src/train.py experiment=panoptic/dales_11g
CUDA内存溢出错误
遇到一些CUDA内存溢出错误 💀💾 ?以下是一些可以调整以减少GPU内存使用的参数,基于错误发生的时机。
影响CUDA内存的参数。
图例: 🟡 预处理 | 🔴 训练 | 🟣 推理(包括训练期间的验证和测试)
参数 | 描述 | 时机 |
---|---|---|
datamodule.xy_tiling | 基于规则XY网格将数据集块分割为xy_tiling^2个较小的块。适用于DALES等方形块。注意这会影响训练步骤数量。 | 🟡🟣 |
datamodule.pc_tiling | 基于主成分将数据集块分割为2^pc_tiling个较小的块。适用于S3DIS和KITTI-360等不同形状的块。注意这会影响训练步骤数量。 | 🟡🟣 |
datamodule.max_num_nodes | 限制训练批次中$P_1$分区节点/超点的数量。 | 🔴 |
datamodule.max_num_edges | 限制训练批次中$P_1$分区边的数量。 | 🔴 |
datamodule.voxel | 增加体素大小将减少预处理、训练和推理时间,但会降低性能。 | 🟡🔴🟣 |
datamodule.pcp_regularization | 分区级别的正则化。越大,超点越少。 | 🟡🔴🟣 |
datamodule.pcp_spatial_weight | 分区中3D位置的重要性。越小,超点越少。 | 🟡🔴🟣 |
datamodule.pcp_cutoff | 最小超点尺寸。越大,超点越少。 | 🟡🔴🟣 |
datamodule.graph_k_max | 超点图中相邻节点的最大数量。越小,超边越少。 | 🟡🔴🟣 |
datamodule.graph_gap | 超点图中相邻超点的最大距离。越小,超边越少。 | 🟡🔴🟣 |
datamodule.graph_chunk | 减小以避免RadiusHorizontalGraph 预处理超点图时出现内存溢出。 | 🟡 |
datamodule.dataloader.batch_size | 控制加载的块数量。每个训练批次由batch_size *datamodule.sample_graph_k 个球形采样组成。在整个验证和测试块上进行推理,无球形采样。 | 🔴🟣 |
datamodule.sample_segment_ratio | 在每个分区级别随机丢弃一部分超点。 | 🔴 |
datamodule.sample_graph_k | 控制训练批次中球形样本的数量。 | 🔴 |
datamodule.sample_graph_r | 控制训练批次中球形样本的半径。设置为sample_graph_r<=0 以使用整个块而不进行球形采样。 | 🔴 |
datamodule.sample_point_min | 控制训练批次中每个超点采样的$P_0$点的最小数量。 | 🔴 |
datamodule.sample_point_max | 控制训练批次中每个超点采样的$P_0$点的最大数量。 | 🔴 |
callbacks.gradient_accumulator.scheduling | 梯度累积。可用于使用较小的批次进行训练,增加训练步骤。 | 🔴 |
💳 致谢
- 本项目使用Lightning-Hydra template构建。
- 本工作的主要数据结构依赖于PyToch Geometric
- 一些点云操作受到Torch-Points3D框架的启发,尽管目前尚未与官方项目合并。
- 对于KITTI-360数据集,使用了官方KITTI-360的一些代码。
- 一些与超点图相关的操作受到Superpoint Graph的启发
- 分层超点分区和图聚类使用Parallel Cut-Pursuit计算
引用我们的工作
如果您的工作使用了本代码的全部或部分,请包含以下引用:
@article{robert2023spt,
title={Efficient 3D Semantic Segmentation with Superpoint Transformer},
author={Robert, Damien and Raguet, Hugo and Landrieu, Loic},
journal={Proceedings of the IEEE/CVF International Conference on Computer Vision},
year={2023}
}
@article{robert2024scalable,
title={Scalable 3D Panoptic Segmentation as Superpoint Graph Clustering},
author={Robert, Damien and Raguet, Hugo and Landrieu, Loic},
journal={Proceedings of the IEEE International Conference on 3D Vision},
year={2024}
}
您可以在arxiv上找到我们的SPT论文 📄和SuperCluster论文 📄。
另外,如果您 ❤️ 或只是使用这个项目,别忘了给这个仓库一个 ⭐,这对我们意义重大!