冬
gcw_IDzXRVNw/UFM-Base-980-ascend
模型介绍
文件和版本
Pull Requests
讨论
分析

UFM-Base-980 昇腾 NPU 适配部署文档

1. 模型简介

UFM-Base-980 (UniFlowMatch) 是卡内基梅隆大学开源的统一密集对应模型的高分辨率版本，能够同时处理光流估计和宽基线匹配任务。

本模型是一种端到端训练的 transformer 模型，支持 980x644 高分辨率推理。

模型架构

属性值
模型类型UFM (统一流匹配)
BackboneDINOv2-Large (24层)
Encoder 隐藏维度1024
Attention Heads16
Patch Size14
推理分辨率980x644
Flow HeadDPT 风格
Uncertainty HeadDPT 风格
Info Sharing全局注意力 (12层)
权重格式safetensors
总权重数613

模型组件

  1. DINOv2 Encoder (336 参数): 24层 ViT-Large 编码器
  2. Flow Head (DPT风格): 从特征图预测光流
  3. Uncertainty Head (DPT风格): 预测不确定性掩码
  4. Info Sharing Module (100 参数): 多视角信息交互

模型特点

  • 高分辨率支持: 原生支持 980x644 分辨率推理
  • 统一框架: 单一模型处理光流和匹配任务
  • DINOv2 Backbone: 基于自监督预训练的视觉 transformer
  • 多任务输出: 同时输出光流和不确定性估计
  • 密集对应: 逐像素预测位移场

输入输出格式

  • 输入: 单张图像 (batch, 3, H, W)，H=980, W=644
  • 输出:
    • Flow: (batch, 2, 980, 644) - UV 光流分量
    • Uncertainty: (batch, 1, 980, 644) - 不确定性掩码

2. 环境依赖清单

2.1 硬件要求

  • 华为昇腾系列 AI 处理器（Ascend 910B/310P 等）
  • 建议 ≥ 3 GB NPU 显存（UFM-Base-980 推理峰值约 2.2 GB）
  • 本次实测使用 910B2（单卡 64 GB）

2.2 软件依赖

依赖项版本要求说明
操作系统CentOS 7.6+ / Ubuntu 18.04+支持主流 Linux 发行版
CANN5.0.RC2+华为昇腾计算架构
Python3.10 / 3.12推荐 3.10+
PyTorch2.0.0+推荐 2.10
torch_npu与 PyTorch 版本匹配昇腾 PyTorch 扩展
safetensors>=0.4.0权重加载

2.3 安装方式

# 1) 激活已预装 torch / torch_npu 的 conda 环境
conda activate pt2100

# 2) 安装模型依赖（使用华为云源）
pip install --index-url https://repo.huaweicloud.com/repository/pypi/simple \
    safetensors torch_npu numpy

# 3) 验证环境
python -c "
import torch, torch_npu
print('torch:', torch.__version__)
print('torch_npu:', torch_npu.__version__)
print('NPU 可用:', torch.npu.is_available())
"

3. 分步推理操作流程

步骤 1：准备工作目录

# 进入适配目录
cd /workspace/agent1/UFM-Base-980-ascend

# 确认上游权重存在
ls -la /workspace/agent2/infinity1096/UFM-Base-980/

步骤 2：检查 NPU 环境

# 检查 NPU 设备状态
npu-smi info

# 验证 torch_npu 安装
python -c "import torch; import torch_npu; print('NPU 可用:', torch.npu.is_available())"

步骤 3：运行推理脚本

# 基本用法：默认参数 (NPU:0, 980x644)
python inference.py --device npu:0

# 指定其他分辨率
python inference.py --device npu:0 --height 980 --width 644

# 指定其他 NPU 卡
python inference.py --device npu:1

# CPU 模式（仅在 NPU 不可用时使用）
python inference.py --device cpu

4. 完整测试用例

测试用例 1：高分辨率光流 + 不确定性估计（NPU:0）

conda activate pt2100
cd /workspace/agent1/UFM-Base-980-ascend
python inference.py --device npu:0

实际输出：

============================================================
UFM-Base-980 昇腾 NPU 推理 (高分辨率)
============================================================

[Device] npu:0
[Model Path] /workspace/agent2/infinity1096/UFM-Base-980

[Loading] Model weights from safetensors...
  Loaded 613 weights

[Creating] Complete UFM model...
  Loaded 336 encoder parameters
  Loaded 18 flow head parameters
  Loaded 18 uncertainty head parameters
  Loaded 100 info_sharing parameters
  Model created successfully

[Model Loaded] Allocated: 1.50 GB, Reserved: 1.66 GB

[npu-smi info]
+------------------------------------------------------------------------------------------------+
| npu-smi 25.5.2                   Version: 25.5.2                                               |
+---------------------------+---------------+----------------------------------------------------+
| NPU   Name                | Health        | Power(W)    Temp(C)           Hugepages-Usage(page)|
| Chip                      | Bus-Id        | AICore(%)   Memory-Usage(MB)  HBM-Usage(MB)        |
+===========================+===============+====================================================+
| 0     910B2               | OK            | 96.6        41                0    / 0             |
| 0                         | 0000:C1:00.0  | 0           0    / 0          5187 / 65536         |
+===========================+===============+====================================================+
+---------------------------+---------------+----------------------------------------------------+
| NPU     Chip              | Process id    | Process name             | Process memory(MB)      |
+===========================+===============+====================================================+
| 0       0                 | 1345197       | python                   | 1818                    |
+===========================+===============+====================================================+


[Creating] Dummy image pair (980x644)...
  Image1 shape: torch.Size([1, 3, 980, 644])
  Image2 shape: torch.Size([1, 3, 980, 644])

[npu-smi info]
+------------------------------------------------------------------------------------------------+
| npu-smi 25.5.2                   Version: 25.5.2                                               |
+---------------------------+---------------+----------------------------------------------------+
| NPU   Name                | Health        | Power(W)    Temp(C)           Hugepages-Usage(page)|
| Chip                      | Bus-Id        | AICore(%)   Memory-Usage(MB)  HBM-Usage(MB)        |
+===========================+===============+====================================================+
| 0     910B2               | OK            | 146.8       41                0    / 0             |
| 0                         | 0000:C1:00.0  | 0           0    / 0          5513 / 65536         |
+===========================+===============+====================================================+


[Flow Results]
  Flow prediction shape: torch.Size([1, 2, 980, 644])
  Flow uv components sample:
    u (horizontal): [1.07, 1.01, 0.95, 0.89, 0.83]
    v (vertical):   [0.44, 0.43, 0.41, 0.40, 0.39]

[Uncertainty Results]
  Uncertainty prediction shape: torch.Size([1, 1, 980, 644])
  Uncertainty sample: [0.098, 0.098, 0.098, 0.098, 0.098]

[Memory After Full Inference] Allocated: 1.48 GB, Reserved: 1.96 GB

[Done]

✅ 验证要点：

  • 高分辨率推理: 原生支持 980x644 分辨率输出
  • 模型加载成功: 加载所有 613 个权重，包括 encoder (336)、flow head (18)、uncertainty head (18)、info_sharing (100)
  • NPU 推理稳定: 光流头和不确定性头在 NPU 上运行推理
  • NPU 进程运行: npu-smi 显示 Python 进程 (PID 1345197) 在 NPU 上运行
  • 光流输出: 输出形状 [1, 2, 980, 644]，UV 分量
  • 不确定性输出: 输出形状 [1, 1, 980, 644]，表示每个像素的不确定性
  • NPU 显存峰值: 约 2.2 GB

5. 常见问题解答

Q1：提示 "NPU 不可用" 怎么办？

解决方案：

  1. 确认已安装 CANN 驱动并设置环境变量（source /usr/local/Ascend/ascend-toolkit/latest/set_env.sh）
  2. 运行 npu-smi info 检查设备状态
  3. 确认 PyTorch 和 torch_npu 版本匹配

Q2：UFM-Base-980 与 UFM-Base 有什么区别？

解释：UFM-Base-980 是高分辨率版本，推理分辨率为 980x644，而 UFM-Base 的推理分辨率为 560x420。980 版本可以处理更精细的对应关系。

Q3：为什么使用更高分辨率？

解释：更高的分辨率可以提供更精确的光流估计，特别是在细节丰富的场景中。UFM-Base-980 特别适合需要高精度对应关系的应用。

6. 文件结构

UFM-Base-980-ascend/
├── inference.py              # 推理脚本（torch_npu 推理）
├── README.md                 # 本文档
├── requirements.txt          # 依赖列表
└── assets/                  # 资源占位目录

上游权重目录（只读，未被修改）：

/workspace/agent2/infinity1096/UFM-Base-980/
├── config.json               # 模型配置
├── model.safetensors         # 权重 (~1.7GB)
└── README.md                 # 原始文档

7. requirements.txt - 依赖列表

# UFM-Base-980 昇腾 NPU 推理环境依赖

# 核心依赖（必须）
torch>=2.0.0
torch_npu
safetensors>=0.4.0
numpy

使用方式：

# 激活环境
conda activate pt2100

# 安装依赖（华为云源）
pip install --index-url https://repo.huaweicloud.com/repository/pypi/simple -r requirements.txt

8. 参考资源

  • UFM 论文 (arXiv:2506.09278)
  • UFM GitHub 仓库
  • UFM HuggingFace
  • DINOv2 论文
  • 华为昇腾文档
  • torch_npu GitHub