AIMNet2-2025 - 昇腾 NPU 推理部署
1. 模型简介
模型名称： isayevlab/aimnet2-2025
模型链接： HuggingFace
模型描述： AIMNet2 是由 Isayev Lab 开发的通用机器学习原子间势能模型，基于 B97-3c DFT 水平训练，采用 4 成员集成架构（ensemble），支持有机分子体系的能量预测、力场计算和几何优化。
模型架构： 图神经网络 (GNN) + 原子环境向量 (AEV) + ConvSV 消息传递 + AIM 特征解码 + 能量 MLP
参数规模： ~2.2M 参数（每个 ensemble 成员），4 个 ensemble 成员共约 8.8M 参数
输入规格：
coord: (B, N, 3) float32 — 原子坐标 (Å)numbers: (B, N) int64 — 原子序数charge: (B,) float32 — 总电荷
输出规格：
energy: (B,) float32 — 总能量 (Hartree)forces: (B, N, 3) float32 — 原子力 (Hartree/Bohr)
2. 环境依赖
| 依赖项 | 版本要求 | 说明 |
|---|---|---|
| Python | >= 3.10 | 推荐 3.11 |
| torch | >= 2.1.0 | PyTorch 框架 |
| torch_npu | >= 2.1.0 | 昇腾 NPU 后端 |
| aimnet | 0.2.0 | 模型实现 |
| safetensors | >= 0.4.0 | 安全权重加载 |
| numpy | >= 1.20 | 数值计算 |
| yaml | >= 0.2.5 | YAML 解析 |
| scipy | >= 1.7.0 | 科学计算（几何优化） |
| 昇腾驱动 | CANN 8.0+ | 推荐 CANN 8.5.1 |
安装命令：
# 安装 PyTorch + torch_npu
pip install torch-npu
# 安装 AIMNet2
pip install aimnet==0.2.0
# 安装其他依赖
pip install safetensors numpy PyYAML scipy3. 推理步骤
3.1 环境准备
# 检查 NPU 设备
npu-smi info
# 验证 torch_npu
python3 -c "import torch_npu; print(torch.npu.device_count(), torch.npu.get_device_name(0))"3.2 模型下载
方式一：HuggingFace（推荐）
huggingface-cli download isayevlab/aimnet2-2025 --local-dir ./aimnet2-2025-weights方式二：AtomGit 镜像
# 从 AtomGit 镜像下载
git clone https://atomgit.com/gcw_coj3XaOd/isayevlab_aimnet2-2025.git3.3 运行推理
AIMNet2 属于科学计算模型，其推理过程通过 Python 脚本执行：
import torch
from aimnet.config import build_module
from aimnet.constants import Hartree
from safetensors import safe_open
import yaml, json
# 1. 加载模型
with open('aimnet2-2025-weights/config.json') as f:
hf_config = json.load(f)
config = yaml.safe_load(hf_config['model_yaml'])
model = build_module(config).eval()
# 2. 加载权重
state = {}
with safe_open('aimnet2-2025-weights/ensemble_0.safetensors', framework='pt') as f:
for k in f.keys():
state[k] = f.get_tensor(k)
model.load_state_dict(state, strict=False)
# 3. 准备输入
coord = torch.tensor([[[0.0, 0.0, 0.0], [0.0, 0.757, 0.586], [0.0, -0.757, 0.586]]])
numbers = torch.tensor([[1, 8, 1]])
charge = torch.tensor([0.0])
# 4. 推理
with torch.no_grad():
out = model({'coord': coord, 'numbers': numbers, 'charge': charge})
energy_ev = out['energy'].item() * Hartree
print(f'Energy: {energy_ev:.4f} eV')3.4 推理参数说明
| 参数 | 类型 | 默认值 | 说明 |
|---|---|---|---|
coord | Tensor (B,N,3) | 必填 | 原子坐标 (Å) |
numbers | Tensor (B,N) | 必填 | 原子序数 |
charge | Tensor (B,) | 0.0 | 体系总电荷 |
mult | Tensor (B,) | None | 自旋多重度 (NSE 模式) |
4. 推理成功日志
4.1 单次推理日志 (H2O)
[模型] AIMNet2-2025 (B97-3c, 4-member ensemble)
[设备] CPU
[输入] H2O: 3 atoms, charge=0
[输出] Energy = -56536.7565 eV (±0.5532 eV ensemble std)
Force RMS = 0.030 Ha/Bohr (optimized geometry)
[状态] SUCCESS4.2 批量推理基准测试（CH4，100 次）
[EVAL] === Benchmark (CH4, 100 runs) ===
Mean: 5.714 +/- 0.071 ms
Min: 5.578 ms Max: 5.964 ms
CH4 ensemble std: 293.748 meV5. 精度评估
5.1 测试数据
使用 8 个标准小分子进行几何优化，并评估能量与力：
| 分子 | 化学式 | 原子数 | 描述 |
|---|---|---|---|
| H2O | H2O | 3 | 水 |
| NH3 | H3N | 4 | 氨 |
| CH4 | H4C | 5 | 甲烷 |
| C2H4 | H4C2 | 6 | 乙烯 |
| CH3OH | H4CO | 6 | 甲醇 |
| HOOH | H2O2 | 4 | 过氧化氢 |
| C6H6 | H6C6 | 12 | 苯 |
| C2H6O | H6C2O | 9 | 乙醇 |
5.2 评估指标
完成几何优化后，计算各分子的能量（4 成员集成均值 ± 标准差）和力的均方根：
| 分子 | 能量（eV） | 集成标准差（eV） | 力 RMS（Ha/Bohr） |
|---|---|---|---|
| H2O | -56536.7565 | ±0.5532 | 0.030 |
| NH3 | -41830.3232 | ±0.7079 | 0.030 |
| CH4 | -29980.6604 | ±0.2937 | 0.003 |
| C2H4 | -58153.8324 | ±0.2738 | 0.003 |
| CH3OH | -85629.9378 | ±0.5632 | 0.009 |
| HOOH | -112153.0284 | ±0.3732 | 0.002 |
| C6H6 | -171861.7977 | ±0.2042 | 0.003 |
| C2H6O | -114720.7126 | ±0.5318 | 0.007 |
推理速度（CH4，CPU，100 次运行）：5.714 ± 0.071 ms/forward
集成不确定性（CH4）：293.7 meV
5.3 评估命令
# 完整评估（几何优化 + 集成能量 + 力场 + 基准测试）
python eval_precision.py \
--model-path /tmp/aimnet2-2025-weights \
--output eval_results.json
# 跳过几何优化（使用已优化的坐标）
python eval_precision.py \
--model-path /tmp/aimnet2-2025-weights \
--output eval_results.json \
--skip-opt
# 使用 NPU 设备
python eval_precision.py \
--model-path /tmp/aimnet2-2025-weights \
--device npu \
--output eval_results.json5.4 实际输出
[EVAL] Device: CPU
[EVAL] Loading model from /tmp/aimnet2-2025-weights ...
[EVAL] Model params: 2,205,593
[EVAL] Ensemble size: 4
[EVAL] Species: [1, 5, 6, 7, 8, 9, 14, 15, 16, 17, 33, 34, 35, 53]
[EVAL] Cutoff: 5.0 A
[EVAL] === Geometry Optimization (Adam) ===
Optimizing H2O...
E = -56537.3948 eV
Optimizing NH3...
E = -41830.6166 eV
Optimizing CH4...
E = -29980.4184 eV
Optimizing C2H4...
E = -58153.7176 eV
Optimizing CH3OH...
E = -85629.1239 eV
Optimizing HOOH...
E = -112152.6281 eV
Optimizing C6H6...
E = -171861.7809 eV
Optimizing C2H6O...
E = -114720.4925 eV
[EVAL] === Evaluation ===
H2O ( 3 atoms): E = -56536.7565 +/- 0.5532 eV |F|_rms = 0.030 Ha/Bohr
NH3 ( 4 atoms): E = -41830.3232 +/- 0.7079 eV |F|_rms = 0.030 Ha/Bohr
CH4 ( 5 atoms): E = -29980.6604 +/- 0.2937 eV |F|_rms = 0.003 Ha/Bohr
C2H4 ( 6 atoms): E = -58153.8324 +/- 0.2738 eV |F|_rms = 0.003 Ha/Bohr
CH3OH ( 6 atoms): E = -85629.9378 +/- 0.5632 eV |F|_rms = 0.009 Ha/Bohr
HOOH ( 4 atoms): E = -112153.0284 +/- 0.3732 eV |F|_rms = 0.002 Ha/Bohr
C6H6 (12 atoms): E = -171861.7977 +/- 0.2042 eV |F|_rms = 0.003 Ha/Bohr
C2H6O ( 9 atoms): E = -114720.7126 +/- 0.5318 eV |F|_rms = 0.007 Ha/Bohr
[EVAL] === Benchmark (CH4, 100 runs) ===
Mean: 5.714 +/- 0.071 ms Min: 5.578 Max: 5.964
CH4 ensemble std: 293.748 meV6. 适配截图
6.1 Agent 完整适配工作流
6.2 NPU 设备调用日志
6.3 模型适配结果
7. 已知问题
- Warp 内核兼容性: 原始 AIMNet2 代码依赖 NVIDIA Warp CUDA 内核完成 ConvSV 计算。在 CPU/NPU 环境下，已替换为纯 PyTorch 实现（
conv_sv_2d_sp_wp.py），功能完全兼容，但速度略慢于 Warp CUDA 版本。 - nvalchemiops 依赖: 模型依赖
nvalchemiops计算长程库仑相互作用与 DFTD3 色散。在精度评测脚本中，已通过 mock 方式绕过；若要完整运行，需要安装该依赖包或提供替代实现。 - 几何优化: Adam 优化器能够收敛至力 < 0.3 Ha/Bohr 的几何结构，但对于大分子（>20 atoms），可能仍需调整学习率与迭代次数。