Pi3 (yyfz233/Pi3) 昇腾 NPU 适配

π³ (Pi-Cubed) 是一个前馈式视觉几何重建网络：输入无序图像集合，无需固定参考视图即可直接预测仿射不变的相机位姿与尺度不变的局部点图（point maps）。得益于完全置换等变（permutation-equivariant）的架构设计，模型对输入顺序天然鲁棒且高度可扩展，在相机位姿估计、单目/视频深度估计与稠密点图估计任务上均达到 SOTA。

本仓库将 yyfz233/Pi3 权重适配到 华为昇腾 Ascend 910 NPU，并提供与 CPU 的精度/性能对比验证。

模型信息

项目	内容
上游模型	yyfz233/Pi3（BSD-2-Clause）
任务	Image-to-3D（相机位姿 + 稠密点云 + 置信度）
架构	DINOv2 ViT-L/14-reg 编码器 + 36 层 RoPE 置换等变 Transformer 解码器（large）+ 点/置信度/相机三个头
参数量	958.7M
输入	`B×N×3×H×W`，[0,1]，H/W 需为 14 的倍数
输出	`points`（全局点云）、`local_points`（局部点图）、`conf`（置信度 logits）、`camera_poses`（4×4 cam-to-world）

适配环境

组件	版本
NPU	Ascend 910 ×2（64GB HBM）
CANN	8.5.1
torch	2.9.0
torch_npu	2.9.0.post1
transformers / Python	4.57.6 / 3.11

适配要点

RoPE2D：上游 cuRoPE2D 为 CUDA 编译扩展，NPU 不可用；自动回退到纯 PyTorch 实现（pos_embed.py 内置 fallback），数值等价。
Attention：未安装 xformers 时 MemEffAttention/Block 自动回退标准 attention；FlashAttentionRope 使用 scaled_dot_product_attention，在 NPU 上走 torch_npu 提供的 SDPA 算子。
autocast：模型 forward 内部硬编码 torch.amp.autocast(device_type='cuda', enabled=False)，该上下文在 NPU 进程下无害，可正常执行；本仓库推理默认 fp32，NPU 亦支持 bf16 混合精度（--precision bf16）。
设备能力判断：example.py 中 torch.cuda.get_device_capability() 在 NPU 环境会报错，本仓库 inference.py 已改为按 --device 参数选择设备与精度。

部署步骤

# 1. 准备目录与权重（权重 ~3.8GB，已放入 models/，不随仓库分发）
mkdir -p /data/Pi3 && cd /data/Pi3
HF_ENDPOINT=https://hf-mirror.com HF_HUB_DOWNLOAD_TIMEOUT=120 python -c \
  "from huggingface_hub import snapshot_download; snapshot_download('yyfz233/Pi3', local_dir='/data/Pi3/models')"

# 2. 克隆上游源码（提供 pi3 包）
git clone https://github.com/yyfz/Pi3.git pi3_repo

# 3. 安装依赖
pip install -r requirements.txt

# 4. 运行推理
python inference.py --device npu   # NPU fp32 推理 + 自动精度对比
python inference.py --device cpu   # CPU fp32 推理 + 自动精度对比
python inference.py --device npu --precision bf16  # NPU bf16 混合精度（可选）

推理示例

import torch, sys
sys.path.insert(0, "pi3_repo")
from pi3.models.pi3 import Pi3

device = "npu" if torch.npu.is_available() else "cpu"
model = Pi3.from_pretrained("models").to(device).eval()

imgs = torch.rand(1, 2, 3, 392, 392, device=device)  # (B,N,3,H,W), [0,1]
with torch.no_grad():
    res = model(imgs)

res["points"].shape         # (1, 2, 392, 392, 3) 全局点云
res["camera_poses"].shape   # (1, 2, 4, 4) cam-to-world

CPU vs NPU 精度与性能对比

测试条件：2 视角 392×392 合成双目输入（PIL 合成，含视差），fp32，各运行一次完整前向。

输出	cos_sim	max_abs_err	mean_abs_err
points	0.999988	0.064541	0.003406
local_points	0.999995	0.066114	0.002975
conf	0.999997	0.142471	0.022954
camera_poses	0.999999	0.001818	0.000505

指标	CPU (fp32)	NPU (fp32)	NPU 加速比
推理耗时	9.478 s	4.917 s	1.93×

所有输出 cos_sim ≥ 0.9999，mean_abs_err 处于 fp32 舍入与算子实现差异的正常范围，NPU 与 CPU 精度对齐。

仓库结构

├── inference.py        # NPU/CPU 推理 + 精度对比（--device npu|cpu）
├── requirements.txt    # Python 依赖
├── readme.md           # 本文件
├── assets/
│   ├── agent_workflow.png   # 适配工作流
│   ├── npu_device_call.png  # NPU 设备与推理证据
│   └── model_result.png     # 模型结果信息图
├── outputs/            # 推理输出（npz/json，gitignore）
├── models/             # 权重目录（gitignore）
└── pi3_repo/           # 上游源码（gitignore）

引用

@misc{wang2025pi3,
      title={π³: Scalable Permutation-Equivariant Visual Geometry Learning},
      author={Yifan Wang and Jianjun Zhou and Haoyi Zhu and Wenzheng Chang and Yang Zhou and Zizun Li and Junyi Chen and Jiangmiao Pang and Chunhua Shen and Tong He},
      year={2025},
      eprint={2507.13347},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}