Sonata

简而言之： 本仓库提供了用于3D点云下游任务的自监督预训练Point Transformer V3模型。

本仓库是论文**Sonata: 自监督学习可靠的点表示的官方项目仓库，主要用于提供预训练模型、推理代码和可视化演示。若需复现Sonata的预训练过程，请参考我们的Pointcept**代码库。 [ 预训练 ] [ Sonata ] - [ 主页 ] [ 论文 ] [ 引用 ]

亮点

• 2025年4月 🚀：Sonata 被选为CVPR 2025的亮点报告之一（仅占提交论文的3.0%）！
• 2025年3月：Sonata 被CVPR 2025接收！我们随**Pointcept** v1.6.0一同发布了预训练代码，并在本仓库中提供了易于使用的推理演示和可视化工具，搭配我们的预训练模型权重。强烈建议用户从本仓库开始进行**快速入门**。

概览

• 安装
• 快速入门
• 引用

安装

本仓库提供了两种安装方式：独立模式和包模式。

• 对于希望快速进行推理和可视化操作的用户，推荐使用独立模式。我们通过conda环境文件提供了一种最简便的环境搭建方式。只需运行以下命令，即可轻松安装包含cuda和pytorch在内的完整环境：

  # 创建并激活名为 'sonata' 的 conda 环境
  # cuda: 12.4, pytorch: 2.5.0
  
  # 如果本地已安装 cuda，请先执行 `unset CUDA_PATH`
  conda env create -f environment.yml --verbose
  conda activate sonata
  

  我们默认安装了*FlashAttention，但并非必需。若本地环境中无法使用 FlashAttention，也无需担心；请参阅【快速入门】中的“模型”部分获取解决方案。*

• 对于希望将我们的模型集成到自身代码库中的用户，推荐使用包模式。我们提供了setup.py文件用于安装。您可以通过以下命令完成安装：

  # 确保本地已安装 Cuda 和 Pytorch
  
  # CUDA_VERSION：本地环境的 CUDA 版本（例如 124），可通过运行 'nvcc --version' 查看
  # TORCH_VERSION：本地环境的 PyTorch 版本（例如 2.5.0），可通过运行 'python -c "import torch; print(torch.__version__)"' 查看
  pip install spconv-cu${CUDA_VERSION}
  pip install torch-scatter -f https://data.pyg.org/whl/torch-{TORCH_VERSION}+cu${CUDA_VERSION}.html
  pip install git+https://github.com/Dao-AILab/flash-attention.git
  pip install huggingface_hub timm
  
  # （可选，或直接将 sonata 文件夹复制到您的项目中）
  python setup.py install
  

  此外，为了运行我们的演示代码，还需要安装以下包：

  pip install open3d fast_pytorch_kmeans psutil numpy==1.26.4  # 目前，open3d 尚不支持 numpy 2.x
  

快速入门

让我们首先使用我们的预训练 PTv3 模型 Sonata 进行一些简单的可视化演示：

• 可视化。 我们在 demo 文件夹中提供了相似性热图和 PCA 可视化演示。你可以运行以下命令来可视化结果：

  export PYTHONPATH=./
  python demo/0_pca.py
  python demo/1_similarity.py
  python demo/2_sem_seg.py  # 在 ScanNet 数据集上使用线性探测头
  

接下来，以下是使用我们的 Sonata 对自定义数据进行推理的说明：

• 数据。 将你的数据组织成一个字典，格式如下：

  # 单个点云
  point = {
    "coord": numpy.array,  # (N, 3)
    "color": numpy.array,  # (N, 3)
    "normal": numpy.array,  # (N, 3)
    "segment": numpy.array,  # (N,) 可选
  }
  
  # 批量点云
  
  # 从这里查看批量点云的数据结构：
  # https://github.com/Pointcept/Pointcept#offset
  point = {
    "coord": numpy.array,  # (N, 3)
    "color": numpy.array,  # (N, 3)
    "normal": numpy.array,  # (N, 3)
    "batch": numpy.array,  # (N,) 可选
    "segment": numpy.array,  # (N,) 可选
  }
  

  你可以通过运行以下命令加载一个示例数据：

  point = sonata.data.load("sample1")
  

• 变换。 数据变换流程与 Pointcept 代码库中使用的相同。你可以使用以下代码构建变换流程：

  config = [
      dict(type="CenterShift", apply_z=True),
      dict(
          type="GridSample",
          grid_size=0.02,
          hash_type="fnv",
          mode="train",
          return_grid_coord=True,
          return_inverse=True,
      ),
      dict(type="NormalizeColor"),
      dict(type="ToTensor"),
      dict(
          type="Collect",
          keys=("coord", "grid_coord", "color", "inverse"),
          feat_keys=("coord", "color", "normal"),
      ),
  ]
  transform = sonata.transform.Compose(config)
  

  你也可以通过运行以下命令获取上述默认的推理增强流程：

  transform = sonata.transform.default()
  

• 模型。 通过运行以下命令加载预训练模型：

  # 从 Huggingface 加载预训练模型
  # 支持的模型: "sonata"
  # 检查点缓存在 ~/.cache/sonata/ckpt 中，可以通过设置 'download_root' 自定义路径
  model = sonata.model.load("sonata", repo_id="facebook/sonata").cuda()
  
  # 或者
  from sonata.model import PointTransformerV3
  model = PointTransformerV3.from_pretrained("facebook/sonata").cuda()
  
  # 从本地路径加载预训练模型
  # 假设检查点文件存储在 'ckpt' 文件夹中
  model = sonata.model.load("ckpt/sonata.pth").cuda()
  
  # 检查点文件包含预训练模型的配置和状态字典
  

  如果 FlashAttention 不可用，可以使用以下代码加载预训练模型：

  custom_config = dict(
      enc_patch_size=[1024 for _ in range(5)],
      enable_flash=False,  # 如有必要，可减小补丁大小
  )
  model = sonata.load("sonata", repo_id="facebook/sonata", custom_config=custom_config).cuda()
  # 或者
  from sonata.model import PointTransformerV3
  model = PointTransformerV3.from_pretrained("facebook/sonata", **custom_config).cuda()
  

• 推理。 通过运行以下命令进行推理：

  point = transform(point)
  for key in point.keys():
      if isinstance(point[key], torch.Tensor):
          point[key] = point[key].cuda(non_blocking=True)
  point = model(point)
  

  由于 Sonata 是一个预训练的 仅编码器 PTv3 模型，模型的默认输出是经过层次化编码后的点云。可以使用以下代码将编码后的点特征映射回原始尺度：

  for _ in range(2):
      assert "pooling_parent" in point.keys()
      assert "pooling_inverse" in point.keys()
      parent = point.pop("pooling_parent")
      inverse = point.pop("pooling_inverse")
      parent.feat = torch.cat([parent.feat, point.feat[inverse]], dim=-1)
      point = parent
  while "pooling_parent" in point.keys():
      assert "pooling_inverse" in point.keys()
      parent = point.pop("pooling_parent")
      inverse = point.pop("pooling_inverse")
      parent.feat = point.feat[inverse]
      point = parent
  

  然而，在数据转换过程中，我们执行了 GridSampling 操作，这会导致输入网络的点数与原始点云不匹配。可以使用以下代码进一步将特征映射回原始点云：

  feat = point.feat[point.inverse]
  

引用

如果你发现 Sonata 对你的研究有帮助，请考虑引用我们的工作以表示感谢。(੭ˊ꒳​ˋ)੭✧

@inproceedings{wu2025sonata,
    title={Sonata: 自监督学习可靠的点表示},
    author={Wu, Xiaoyang 和 DeTone, Daniel 和 Frost, Duncan 和 Shen, Tianwei 和 Xie, Chris 和 Yang, Nan 和 Engel, Jakob 和 Newcombe, Richard 和 Zhao, Hengshuang 和 Straub, Julian},
    booktitle={CVPR},
    year={2025}
}

@inproceedings{wu2024ptv3,
    title={Point Transformer V3: 更简单、更快、更强},
    author={Wu, Xiaoyang 和 Jiang, Li 和 Wang, Peng-Shuai 和 Liu, Zhijian 和 Liu, Xihui 和 Qiao, Yu 和 Ouyang, Wanli 和 He, Tong 和 Zhao, Hengshuang},
    booktitle={CVPR},
    year={2024}
}

@inproceedings{wu2024ppt,
    title={通过多数据集点提示训练实现大规模 3D 表征学习},
    author={Wu, Xiaoyang 和 Tian, Zhuotao 和 Wen, Xin 和 Peng, Bohao 和 Liu, Xihui 和 Yu, Kaicheng 和 Zhao, Hengshuang},
    booktitle={CVPR},
    year={2024}
}

@inproceedings{wu2023masked,
  title={遮蔽场景对比：一种可扩展的无监督 3D 表征学习框架},
  author={Wu, Xiaoyang 和 Wen, Xin 和 Liu, Xihui 和 Zhao, Hengshuang},
  journal={CVPR},
  year={2023}
}

@inproceedings{wu2022ptv2,
    title={Point transformer V2: 分组向量注意力和基于分区的池化},
    author={Wu, Xiaoyang 和 Lao, Yixing 和 Jiang, Li 和 Liu, Xihui 和 Zhao, Hengshuang},
    booktitle={NeurIPS},
    year={2022}
}

@misc{pointcept2023,
    title={Pointcept: 用于点云感知研究的代码库},
    author={Pointcept 贡献者},
    howpublished={\url{https://github.com/Pointcept/Pointcept}},
    year={2023}
}

如何贡献

我们欢迎各种贡献！请参阅 CONTRIBUTING 和我们的 行为准则，了解如何开始。

许可证

• Sonata 代码由 Meta 在 Apache 2.0 许可证 下发布；
• Sonata 模型权重在 CC-BY-NC 4.0 许可证 下发布（受 HM3D、ArkitScenes 等数据集的 NC 条款限制）。

Sonata 快速上手指南

Sonata 是一个基于 Point Transformer V3 (PTv3) 的自监督预训练模型，专为 3D 点云下游任务（如语义分割）设计。本指南将帮助你快速配置环境并运行推理与可视化演示。

环境准备

在开始之前，请确保你的系统满足以下要求：

• 操作系统: Linux (推荐 Ubuntu)
• GPU: 支持 CUDA 的 NVIDIA 显卡
• CUDA 版本: 推荐 12.4 (其他版本需在包模式下手动指定)
• Python: 3.8+
• 包管理器: Conda (推荐用于独立模式)

安装步骤

本项目提供两种安装模式：独立模式（推荐用于快速推理和可视化）和 包模式（用于集成到现有项目）。

方式一：独立模式（推荐）

此模式通过 Conda 一键安装包含 CUDA 和 PyTorch 的完整环境，最适合新手快速体验。

# 如果本地已安装 CUDA，建议先执行 unset CUDA_PATH 以避免冲突
unset CUDA_PATH

# 创建并激活名为 'sonata' 的 conda 环境 (自动安装 CUDA 12.4, PyTorch 2.5.0)
conda env create -f environment.yml --verbose
conda activate sonata

注意: 默认会尝试安装 FlashAttention 以加速。如果安装失败或不具备相应硬件，代码会自动降级处理，不影响基本运行。

方式二：包模式

如果你希望将 Sonata 集成到已有的深度学习环境中，请使用此模式。请确保本地已正确安装 CUDA 和 PyTorch。

# 获取本地环境版本
# CUDA_VERSION: 例如 124 (通过 nvcc --version 查看)
# TORCH_VERSION: 例如 2.5.0 (通过 python -c "import torch; print(torch.__version__)" 查看)

# 安装核心依赖 (请替换 ${CUDA_VERSION} 和 ${TORCH_VERSION} 为实际值)
pip install spconv-cu${CUDA_VERSION}
pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH_VERSION}+cu${CUDA_VERSION}.html
pip install git+https://github.com/Dao-AILab/flash-attention.git
pip install huggingface_hub timm

# 安装 Sonata 包
python setup.py install

# 若需运行演示代码 (demo)，还需安装以下额外依赖
pip install open3d fast_pytorch_kmeans psutil numpy==1.26.4

基本使用

1. 运行可视化演示

安装完成后，你可以直接运行提供的脚本来查看预训练模型的效果（包括 PCA 降维可视化、相似度热力图和语义分割结果）。

export PYTHONPATH=./

# PCA 可视化
python demo/0_pca.py

# 相似度热力图
python demo/1_similarity.py

# ScanNet 数据集上的线性探测语义分割
python demo/2_sem_seg.py

2. 自定义数据推理

以下是使用 Sonata 对自定义点云数据进行推理的最小化代码示例。

第一步：准备数据

数据需组织为字典格式，包含坐标、颜色等信息。

import numpy as np
import sonata

# 示例：加载内置样本或构建自己的数据
# point = sonata.data.load("sample1") 

# 自定义数据结构示例
point = {
    "coord": np.random.rand(1000, 3).astype(np.float32),  # (N, 3) 坐标
    "color": np.random.rand(1000, 3).astype(np.float32),  # (N, 3) 颜色
    "normal": np.random.rand(1000, 3).astype(np.float32), # (N, 3) 法向量
}

第二步：配置变换管道

使用默认的推理增强流程。

# 获取默认变换配置
transform = sonata.transform.default()

# 应用变换
point = transform(point)

第三步：加载模型

从 Hugging Face 自动下载预训练权重，或从本地加载。

import torch
import sonata

# 方案 A: 从 Hugging Face 加载 (自动缓存至 ~/.cache/sonata/ckpt)
model = sonata.model.load("sonata", repo_id="facebook/sonata").cuda()

# 方案 B: 如果未安装 FlashAttention，需禁用它并调整 patch_size
# custom_config = dict(enc_patch_size=[1024 for _ in range(5)], enable_flash=False)
# model = sonata.model.load("sonata", repo_id="facebook/sonata", custom_config=custom_config).cuda()

第四步：执行推理与特征还原

模型输出为分层编码后的特征，需要映射回原始点云尺度。

# 将数据移至 GPU
for key in point.keys():
    if isinstance(point[key], torch.Tensor):
        point[key] = point[key].cuda(non_blocking=True)

# 前向传播
point = model(point)

# 步骤 1: 逆向池化操作，将特征逐层还原
for _ in range(2):
    assert "pooling_parent" in point.keys()
    parent = point.pop("pooling_parent")
    inverse = point.pop("pooling_inverse")
    parent.feat = torch.cat([parent.feat, point.feat[inverse]], dim=-1)
    point = parent

while "pooling_parent" in point.keys():
    parent = point.pop("pooling_parent")
    inverse = point.pop("pooling_inverse")
    parent.feat = point.feat[inverse]
    point = parent

# 步骤 2: 通过 GridSampling 的逆映射，将特征对齐到原始输入点数
final_features = point.feat[point.inverse]

print(f"原始点数：{point['coord'].shape[0]}, 输出特征维度：{final_features.shape[-1]}")