ML
安装 | 快速开始 | 结构 | 任务与算法 | 模型库 | 数据集 | 使用方法 | 贡献
Open3D-ML 是 Open3D 的扩展,用于 3D 机器学习任务。它建立在 Open3D 核心库之上,并通过机器学习工具扩展了 3D 数据处理功能。该仓库侧重于诸如语义点云分割等应用,并提供了可应用于常见任务的预训练模型以及用于训练的管道。
Open3D-ML 支持 TensorFlow 和 PyTorch,可轻松集成到现有项目中,并提供了独立于 ML 框架的通用功能,如数据可视化。
安装
用户
Open3D-ML 已集成在 Open3D v0.11+ 的 Python 发行版中,并与以下版本的 ML 框架兼容。
- PyTorch 2.0.*
- TensorFlow 2.13.*(macOS,Linux 见下文)
- CUDA 10.1, 11.*(在
GNU/Linux x86_64
上可选)
你可以使用以下命令安装 Open3D:
# 确保你拥有最新版本的 pip
pip install --upgrade pip
# 安装 open3d
pip install open3d
要安装与 PyTorch 或 TensorFlow 兼容的版本,你可以使用相应的需求文件:
# 安装与 TensorFlow 兼容的版本
pip install -r requirements-tensorflow.txt
# 安装与 PyTorch 兼容的版本
pip install -r requirements-torch.txt
# 在 Linux 上安装支持 CUDA 的 PyTorch 版本
pip install -r requirements-torch-cuda.txt
要测试安装是否成功,使用以下命令:
# 使用 PyTorch
$ python -c "import open3d.ml.torch as ml3d"
# 或使用 TensorFlow
$ python -c "import open3d.ml.tf as ml3d"
如果你需要使用不同版本的 ML 框架或 CUDA,我们建议你 从源码构建 Open3D 或 在 Docker 中构建 Open3D。
从 v0.18 起,在 Linux 上 PyPI Open3D 包不再原生支持 TensorFlow,因为 PyTorch 与 TensorFlow 之间的构建不兼容 [详情参见 Python 3.11 支持 PR]。如果你想在 Linux 上使用带有 TensorFlow 支持的 Open3D,你可以 在 Docker 中从源码构建带有 TensorFlow 支持的 Open3D 包 (但不支持 PyTorch),命令如下:
cd docker
# 为 Python 3.10 构建带有 TensorFlow 支持的 open3d 和 open3d-cpu 包
export BUILD_PYTORCH_OPS=OFF BUILD_TENSORFLOW_OPS=ON
./docker_build.sh cuda_wheel_py310
快速开始
读取数据集
数据集命名空间包含了读取常见数据集的类。这里我们读取 SemanticKITTI 数据集并将其可视化。
import open3d.ml.torch as ml3d # 或者 open3d.ml.tf as ml3d
# 通过指定 dataset_path 来构建数据集
dataset = ml3d.datasets.SemanticKITTI(dataset_path='/path/to/SemanticKITTI/')
# 获取将训练集、验证集和测试集合并的 “all” 分割
all_split = dataset.get_split('all')
# 打印第一个数据的属性
print(all_split.get_attr(0))
# 打印第一个点云的形状
print(all_split.get_data(0)['point'].shape)
# 使用可视化工具显示前 100 帧
vis = ml3d.vis.Visualizer()
vis.visualize_dataset(dataset, 'all', indices=range(100))
加载配置文件
模型、数据集和管道的配置保存在 ml3d/configs
目录下。用户也可以构建自己的 yaml 文件来记录自定义配置。以下是读取配置文件并据此构建模块的示例。
import open3d.ml as _ml3d
import open3d.ml.torch as ml3d # 或者 open3d.ml.tf as ml3d
framework = "torch" # 或 tf
cfg_file = "ml3d/configs/randlanet_semantickitti.yml"
cfg = _ml3d.utils.Config.load_from_file(cfg_file)
# 按名称获取类
Pipeline = _ml3d.utils.get_module("pipeline", cfg.pipeline.name, framework)
Model = _ml3d.utils.get_module("model", cfg.model.name, framework)
Dataset = _ml3d.utils.get_module("dataset", cfg.dataset.name)
# 使用配置文件中的参数构建实例
cfg.dataset['dataset_path'] = "/path/to/your/dataset"
dataset = Dataset(cfg.dataset.pop('dataset_path', None), **cfg.dataset)
model = Model(**cfg.model)
pipeline = Pipeline(model, dataset, **cfg.pipeline)
语义分割
运行预训练模型进行语义分割
基于之前的示例,我们可以实例化一个带有预训练模型的管道,并在数据集的点云上运行它。关于如何获取预训练模型的权重,请参见 模型库。
import os
import open3d.ml as _ml3d
import open3d.ml.torch as ml3d
cfg_file = "ml3d/configs/randlanet_semantickitti.yml"
cfg = _ml3d.utils.Config.load_from_file(cfg_file)
model = ml3d.models.RandLANet(**cfg.model)
cfg.dataset['dataset_path'] = "/path/to/your/dataset"
dataset = ml3d.datasets.SemanticKITTI(cfg.dataset.pop('dataset_path', None), **cfg.dataset)
pipeline = ml3d.pipelines.SemanticSegmentation(model, dataset=dataset, device="gpu", **cfg.pipeline)
# 下载权重文件
ckpt_folder = "./logs/"
os.makedirs(ckpt_folder, exist_ok=True)
ckpt_path = ckpt_folder + "randlanet_semantickitti_202201071330utc.pth"
randlanet_url = "https://storage.googleapis.com/open3d-releases/model-zoo/randlanet_semantickitti_202201071330utc.pth"
if not os.path.exists(ckpt_path):
cmd = "wget {} -O {}".format(randlanet_url, ckpt_path)
os.system(cmd)
# 加载参数
pipeline.load_ckpt(ckpt_path=ckpt_path)
test_split = dataset.get_split("test")
data = test_split.get_data(0)
# 对单个样本运行推理
# 返回包含 “predict_labels” 和 “predict_scores” 的字典
result = pipeline.run_inference(data)
# 在测试集上评估性能;这将把日志写入 ‘./logs’
pipeline.run_test()
用户也可以 使用预定义的脚本 来加载预训练权重并运行测试。
训练语义分割模型
与推理类似,管道提供了在数据集上训练模型的接口。
# 使用缓存来存储预处理结果(默认路径为 './logs/cache')
dataset = ml3
<img width="640" src="https://yellow-cdn.veclightyear.com/2b54e442/5add5afd-0c92-46b3-807f-5272c3577226.jpg">
### 3D物体检测
#### 运行预训练模型进行3D物体检测
3D物体检测模型类似于语义分割模型。我们可以用预训练模型实例化一个管道,并在我们的数据集的点云上运行它。查看[模型库](#model-zoo)以获取预训练模型的权重。
```python
import os
import open3d.ml as _ml3d
import open3d.ml.torch as ml3d
cfg_file = "ml3d/configs/pointpillars_kitti.yml"
cfg = _ml3d.utils.Config.load_from_file(cfg_file)
model = ml3d.models.PointPillars(**cfg.model)
cfg.dataset['dataset_path'] = "/path/to/your/dataset"
dataset = ml3d.datasets.KITTI(cfg.dataset.pop('dataset_path', None), **cfg.dataset)
pipeline = ml3d.pipelines.ObjectDetection(model, dataset=dataset, device="gpu", **cfg.pipeline)
# 下载权重。
ckpt_folder = "./logs/"
os.makedirs(ckpt_folder, exist_ok=True)
ckpt_path = ckpt_folder + "pointpillars_kitti_202012221652utc.pth"
pointpillar_url = "https://storage.googleapis.com/open3d-releases/model-zoo/pointpillars_kitti_202012221652utc.pth"
if not os.path.exists(ckpt_path):
cmd = "wget {} -O {}".format(pointpillar_url, ckpt_path)
os.system(cmd)
# 加载参数。
pipeline.load_ckpt(ckpt_path=ckpt_path)
test_split = dataset.get_split("test")
data = test_split.get_data(0)
# 在单个示例上运行推理。
# 返回包含'predict_labels'和'predict_scores'的字典。
result = pipeline.run_inference(data)
# 在测试集上评估性能; 这将写日志到'./logs'。
pipeline.run_test()
用户还可以使用预定义脚本加载预训练权重并运行测试。
训练3D物体检测模型
与推理类似,管道为在数据集上训练模型提供了接口。
# 使用缓存存储预处理的结果(默认路径为'./logs/cache')
dataset = ml3d.datasets.KITTI(dataset_path='/path/to/KITTI/', use_cache=True)
# 用随机初始化创建模型。
model = PointPillars()
pipeline = ObjectDetection(model=model, dataset=dataset, max_epoch=100)
# 在控制台中打印训练进度。
pipeline.run_train()
下面是一个使用KITTI进行可视化的示例。示例显示了对KITTI数据集使用边界框。
欲了解更多示例,请参见examples/
和scripts/
目录。您还可以在配置文件中启用保存训练摘要,并使用tensorboard可视化真实数据和结果。有关详细信息,请参见本教程。
使用预定义脚本
scripts/run_pipeline.py
提供了一个简单的接口,用于在数据集上训练和评估模型。它省去了定义特定模型和传递精确配置的麻烦。
python scripts/run_pipeline.py {tf/torch} -c <path-to-config> --pipeline {SemanticSegmentation/ObjectDetection} --<extra args>
该脚本可用于语义分割和物体检测。你必须在pipeline
参数中指定SemanticSegmentation或ObjectDetection。请注意,extra args
将优先于配置文件中存在的相同参数。因此,与其在配置文件中更改参数,不如在启动脚本时将其作为命令行参数传递。
例如:
# 使用torch在SemanticKITTI上启动RandLANet的训练。
python scripts/run_pipeline.py torch -c ml3d/configs/randlanet_semantickitti.yml --dataset.dataset_path <path-to-dataset> --pipeline SemanticSegmentation --dataset.use_cache True
# 使用torch在KITTI上启动PointPillars的测试。
python scripts/run_pipeline.py torch -c ml3d/configs/pointpillars_kitti.yml --split test --dataset.dataset_path <path-to-dataset> --pipeline ObjectDetection --dataset.use_cache True
要获得进一步帮助,运行python scripts/run_pipeline.py --help
。
仓库结构
Open3D-ML的核心部分位于ml3d
子文件夹中,集成在Open3D的ml
命名空间中。除了核心部分,examples
和scripts
目录提供了支持脚本,以帮助建立训练管道或在数据集上运行网络。
├─ docs # 文档的Markdown和rst文件
├─ examples # 示例脚本和笔记本的位置
├─ ml3d # 集成到open3d中的包根目录
├─ configs # 模型配置文件
├─ datasets # 通用数据集代码; 将集成为open3d.ml.{tf,torch}.datasets
├─ metrics # 用于评估ML模型的指标
├─ utils # 框架无关的实用工具; 将作为open3d.ml.{tf,torch}.utils提供
├─ vis # 特定于ML的可视化功能
├─ tf # TensorFlow特定代码的目录,与ml3d/torch结构相同。
│ # 这将作为open3d.ml.tf提供
├─ torch # PyTorch特定代码的目录;作为open3d.ml.torch提供
├─ dataloaders # 框架特定的数据集代码,例如可以使用
│ # 通用数据集代码的包装器
├─ models # 模型代码
├─ modules # 较小的模块,例如指标和损失
├─ pipelines # 用于任务(如语义分割)的管道
├─ utils # 实用工具 <>
├─ scripts # 训练演示脚本和数据集下载脚本
任务和算法
语义分割
对于语义分割任务,我们使用包含所有类的平均交并比(mIoU)来衡量不同方法的性能。 该表显示了可用于分割任务的模型和数据集及其各自的得分。每个分数链接到各自的权重文件。 <SOURCE_TEXT>
模型 / 数据集 | SemanticKITTI | Toronto 3D | S3DIS | Semantic3D | Paris-Lille3D | ScanNet |
---|---|---|---|---|---|---|
RandLA-Net (tf) | 53.7 | 73.7 | 70.9 | 76.0 | 70.0* | - |
RandLA-Net (torch) | 52.8 | 74.0 | 70.9 | 76.0 | 70.0* | - |
KPConv (tf) | 58.7 | 65.6 | 65.0 | - | 76.7 | - |
KPConv (torch) | 58.0 | 65.6 | 60.0 | - | 76.7 | - |
SparseConvUnet (torch) | - | - | - | - | - | 68 |
SparseConvUnet (tf) | - | - | - | - | - | 68.2 |
PointTransformer (torch) | - | - | 69.2 | - | - | - |
PointTransformer (tf) | - | - | 69.2 | - | - | - |
(*) 使用原作者提供的权重。
物体检测
在物体检测任务中,我们使用鸟瞰图 (BEV) 和 3D 的平均精度均值 (mAP) 来衡量不同方法的性能。表格显示了物体检测任务中可用的模型和数据集及其相应的分数。每个分数链接到相应的权重文件。根据 KITTI 的验证标准,这些模型在验证子集上进行了评估。模型针对三类(汽车、行人和自行车骑行者)进行了训练。所计算的值是所有类别在所有难度级别上的 mAP 的平均值。对于 Waymo 数据集,模型针对三类(行人、车辆、自行车骑行者)进行了训练。
模型 / 数据集 | KITTI [BEV / 3D] @ 0.70 | Waymo (BEV / 3D) @ 0.50 |
---|---|---|
PointPillars (tf) | 61.6 / 55.2 | - |
PointPillars (torch) | 61.2 / 52.8 | avg: 61.01 / 48.30 | best: 61.47 / 57.55 1 |
PointRCNN (tf) | 78.2 / 65.9 | - |
PointRCNN (torch) | 78.2 / 65.9 | - |
PointRCNN 训练
要使用用于训练的真实数据采样数据增强,我们可以按如下方式生成真实数据数据库:
python scripts/collect_bboxes.py --dataset_path <path_to_data_root>
这将生成一个由训练集拆分出的对象组成的数据库。建议对 KITTI 等对象稀疏的数据集使用此增强。
PointRCNN 的两个阶段分别进行训练。要使用 PyTorch 训练 PointRCNN 的提案生成阶段,请运行以下命令:
# 训练 RPN 100 个 epoch。
python scripts/run_pipeline.py torch -c ml3d/configs/pointrcnn_kitti.yml --dataset.dataset_path <path-to-dataset> --mode RPN --epochs 100
在获得训练良好的 RPN 网络后,我们可以在冻结 RPN 权重的情况下训练 RCNN 网络。
# 训练 RCNN 70 个 epoch。
python scripts/run_pipeline.py torch -c ml3d/configs/pointrcnn_kitti.yml --dataset.dataset_path <path-to-dataset> --mode RCNN --model.ckpt_path <path_to_checkpoint> --epochs 100
模型库
有关所有权重文件的完整列表,请参见 model_weights.txt 和 MD5 校验文件 model_weights.md5。
数据集
以下是我们提供的数据集读取类的数据集列表。
- SemanticKITTI (项目页面)
- Toronto 3D (github)
- Semantic 3D (项目页面)
- S3DIS (项目页面)
- Paris-Lille 3D (项目页面)
- Argoverse (项目页面)
- KITTI (项目页面)
- Lyft (项目页面)
- nuScenes (项目页面)
- Waymo (项目页面)
- ScanNet (项目页面)
要下载这些数据集,请访问相应的网页并查看 [scripts/download_datasets
](https://github.com/isl-org/Open3D-ML/tree/main/scripts/download_d
贡献
有许多方式可以为该项目做出贡献。你可以:
- 实现一个新模型
- 添加读取新数据集的代码
- 分享现有模型的参数和权重
- 报告问题和错误
请将你的拉取请求提交到 dev 分支。 Open3D 是一个社区合作项目。我们欢迎并庆祝来自社区的贡献!
如果你想分享你训练的模型的权重文件,请在拉取请求中附上或链接该权重文件。 对于错误和问题,请提交问题。 请查看我们的沟通渠道以便与社区取得联系。
沟通渠道
- 论坛:讨论Open3D的使用。
- Discord 聊天:在线聊天、讨论以及与其他用户和开发者的合作。
引用
如果你使用Open3D,请引用我们的工作 (pdf)。
@article{Zhou2018,
author = {Qian-Yi Zhou and Jaesik Park and Vladlen Koltun},
title = {{Open3D}: {A} Modern Library for {3D} Data Processing},
journal = {arXiv:1801.09847},
year = {2018},
}
Footnotes
-
avg. 指标是 4、8、16 和 32 GPUs 进行的三组训练的平均值。训练在 30 个 epoch 后停止。最佳训练运行的模型检查点可用。 ↩