REANN模型在昇腾的部署实践

weixin_45870772

115人浏览 · 2026-11-09 14:38:55

weixin_45870772 · 2026-11-09 14:38:55 发布

作者：昇腾实战派

知识地图：https://blog.csdn.net/Lumos_Lovegood/article/details/161601003

背景概述

构建精确且高效的原子间势能函数（机器学习势能，MLP）是计算化学和材料科学的核心目标之一。传统的原子描述符方法（如对称函数、Coulomb矩阵等）在表达能力和系统可扩展性上存在局限。如何设计一种既能准确捕获多体相互作用、又能系统性地改进的原子表示方法，是该领域的关键挑战。

REANN（Recursively Embedded Atom Neural Network）由中国科学技术大学蒋斌课题组开发，核心论文发表于 Physical Review Letters（2021）。REANN 提出了一种递归嵌入原子密度的全新原子表示方法——通过递归地将多体相关信息嵌入原子密度描述符中，系统性地提升对多体相互作用的描述能力，同时保持旋转不变性和置换等变性。该方法在分子势能面、偶极矩、极化率等多种性质的预测上均达到了高精度。

REANN 是一个基于 PyTorch 的端到端多功能深度神经网络包，支持：

分子间势能（原子力场）训练
偶极矩预测
跃迁偶极矩预测
极化率预测
周期性体系和反应体系

本文介绍 REANN 的昇腾 Ascend NPU 适配版本——利用分布式数据并行在 NPU 上完成高效训练。

模型介绍

REANN 概述

REANN 的核心创新是递归嵌入原子密度（Recursively Embedded Atom Density）。传统嵌入原子方法（EAM）仅使用局部电子密度的标量来描述原子环境，而 REANN 将这一概念推广为递归的多体扩展：

第 0 阶：标量原子密度（类似 EAM，仅包含二体信息）
第 1 阶：向量原子密度（包含三体角度信息）
第 n 阶：张量原子密度（递归嵌入更高阶多体相关）

每一阶的嵌入密度都在前一阶的基础上递归构建，系统性地纳入更高阶的多体相互作用，且通过张量缩并保证旋转不变性。

整体架构

模型组成

径向基函数展开（Gaussian Smearing）
- 使用可学习的高斯中心 rs 和宽度 inta
- 将原子间距离展开为多维径向特征
- 余弦截断函数确保平滑过渡
递归嵌入密度模块（GetDensity）

核心计算流程：
- 计算径向基： $\phi(r) = \exp(\eta (r - r_s)^2) \cdot f_c(r)$
- 对每个原子，按递归阶数 nipsin 构建嵌入密度：
  - 阶数 0：标量聚合（ $\sum_j \phi(r_{ij})$ ）
  - 阶数 1：向量聚合（ $\sum_j \phi(r_{ij}) \cdot \hat{r}_{ij}$ ）
  - 更高阶：通过 opt_einsum 高效计算张量收缩
- 使用可学习超参数 hyper 混合不同径向通道
- 通过 ocmod（消息传递 MLP）更新嵌入
性质预测模块

根据目标性质选择不同的预测模块：
- Property_energy：能量预测（标量，旋转不变）
- Property_force：能量 + 力预测（通过自动微分）
- Property_DM：偶极矩预测（向量，旋转等变）
- Property_POL：极化率预测（二阶张量）
- Property_TDM：跃迁偶极矩预测
全局模型（MODEL）
- 整合嵌入密度和性质预测
- 支持多元素体系
- 支持周期性边界条件

训练流程

信息加载（src.read）→ 初始化（run.train）→ 数据加载（src.dataloader）→ 优化（src.optimize）

异步数据预取：额外线程从 CPU 预取数据到 GPU/NPU
学习率调度：ReduceLROnPlateau，当学习率降到 end_lr 以下时停止
分布式训练：支持 PyTorch DDP 多卡并行

技术特点

旋转不变性：通过张量缩并（Einstein summation）保证
置换等变性：对相同类型原子的置换不变
系统可扩展：递归阶数可任意增加以提升精度
高效实现：使用 opt_einsum 优化张量收缩表达式
多功能：统一框架支持能量、力、偶极矩、极化率等性质

应用场景

分子反应动力学：构建反应势能面
表面催化：气体-表面相互作用（如 CO₂+Ni(100)）
光谱模拟：通过偶极矩/极化率预测红外/拉曼光谱
周期性体系：晶体和表面体系的势能训练
LAMMPS 接口：高效分子动力学模拟

昇腾 NPU 适配版本

迁移动机

REANN 原始基于 PyTorch 实现，支持 CPU 和 CUDA GPU。本项目通过安装 torch_npu 并使用 transfer_to_npu 自动迁移机制，在昇腾 Ascend NPU 上完成分布式训练。

仓库结构

REANN_npu/
├── README.md
├── reann/
│   ├── __main__.py                     # 入口（调用 run.train）
│   ├── __pes__.py
│   ├── run/
│   │   └── train.py                    # 训练主流程
│   ├── src/
│   │   ├── density.py                  # ⭐ 递归嵌入密度核心实现
│   │   ├── MODEL.py                    # 全局模型
│   │   ├── Property_energy.py          # 能量预测
│   │   ├── Property_force.py           # 力预测
│   │   ├── Property_DM.py             # 偶极矩
│   │   ├── Property_POL.py            # 极化率
│   │   ├── Property_TDM.py            # 跃迁偶极矩
│   │   ├── Loss.py                     # 损失函数
│   │   ├── dataloader.py              # 异步数据加载
│   │   ├── optimize.py                # 优化器
│   │   ├── read.py                    # 配置文件解析
│   │   ├── get_neighbour.py           # 邻居列表构建
│   │   ├── cpu_gpu.py                 # 设备管理
│   │   └── EMA.py                     # 指数移动平均
│   ├── ASE/                            # ASE 计算器接口
│   ├── lammps/                         # LAMMPS 接口
│   ├── pes/                            # 势能面预测
│   ├── dm/                             # 偶极矩推理
│   ├── pol/                            # 极化率推理
│   └── inference/                      # 通用推理
├── data/                               # 训练数据
│   └── co2+Ni100/                      # CO₂+Ni(100) 示例
├── example/                            # 示例配置
│   └── co2+ni100/para/
│       ├── input_density               # 嵌入密度参数
│       └── input_nn                    # 网络结构参数
├── manual/                             # 用户手册
└── release_note/

NPU 适配方式

核心文件 src/density.py 顶部添加：

import torch_npu

模型通过 PyTorch DDP 框架的设备管理自动在 NPU 上运行，opt_einsum 的张量收缩表达式由 torch_npu 自动调度到 NPU 执行。

版本信息

软件	版本
HDK	25.0.rc1.1
CANN	8.3.RC1
Python	3.11
PyTorch	2.1.0
torch_npu	2.1.0.post17
opt_einsum	3.2.0
NumPy	1.24.0

环境配置

创建容器

docker run -it -u root \
 --net=host --shm-size=5g \
 --device=/dev/davinci_manager \
 --device=/dev/hisi_hdc \
 --device=/dev/davinci4 \
 --device=/dev/davinci5 \
 -v /usr/local/dcmi:/usr/local/dcmi \
 -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
 -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
 -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
 -v /etc/ascend_install.info:/etc/ascend_install.info \
 --name REANN_test \
 swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-a3-ubuntu22.04-py3.11

镜像获取参考：https://www.hiascend.com/developer/ascendhub/

安装 torch + torch_npu

wget https://gitcode.com/Ascend/pytorch/releases/download/v7.2.0-pytorch2.1.0/torch_npu-2.1.0.post17-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
pip3 install torch_npu-2.1.0.post17-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
pip3 install protobuf==3.20

安装三方依赖

pip3 install numpy==1.24.0
pip3 install opt-einsum==3.2.0

下载代码

git clone https://atomgit.com/AI4Science/REANN.git
cd REANN/reann

测试数据准备

mkdir para && cd para
ln -s ../../example/co2+ni100/para/input_density ./input_density
ln -s ../../example/co2+ni100/para/input_nn ./input_nn
cd ../../data
cp -r co2+Ni100/ co2+ni100/

训练

单卡训练

python3 -m torch.distributed.run \
    --master_addr 127.0.0.1 \
    --master_port 12345 \
    --nproc_per_node=1 \
    --nnodes=1 \
    --standalone ./

多卡训练

python3 -m torch.distributed.run \
    --master_addr 127.0.0.1 \
    --master_port 12345 \
    --nproc_per_node=2 \
    --nnodes=1 \
    --standalone ./

配置文件说明

input_density：控制嵌入密度参数

递归阶数（nipsin）
径向基数量（nwave）
截断半径（cutoff）
邻居原子数量

input_nn：控制网络结构参数

隐藏层维度
网络层数
激活函数
学习率和训练轮次

数据格式

训练数据由 configuration 文件组织，包含：

晶格参数（3×3 矩阵）
周期性边界条件（pbc）
原子类型、质量、坐标、力（可选）
目标性质（能量/偶极矩/极化率）

迁移适配要点

适配策略

组件	适配方式
模型计算（density, property）	`torch_npu` 自动算子调度
`opt_einsum` 张量收缩	PyTorch 后端自动迁移到 NPU
分布式训练（DDP）	通过 `torch.distributed` + HCCL 通信后端
邻居列表构建	CPU 上执行（涉及坐标遍历）
异步数据预取	由 dataloader 额外线程完成 CPU→NPU 传输

已知问题与解决

集群训练 hostname 问题：在 /etc/hosts 中添加 IPv6 域名映射
protobuf 版本冲突：需安装 protobuf==3.20
NumPy 版本：需使用 numpy==1.24.0（兼容性要求）

已知限制

邻居列表构建在 CPU 上执行，大体系可能成为瓶颈
LAMMPS 接口暂不支持 NPU（LAMMPS 使用 LibTorch CPU/CUDA）
ASE 计算器中的 MD 模拟暂需 CPU 运行邻居搜索部分
极大体系（>10000 原子）的 opt_einsum 表达式可能占用较多 NPU 内存

参考文献

Zhang, Y., Hu, C. & Jiang, B. Embedded Atom Neural Network Potentials: Efficient and Accurate Machine Learning with a Physically Inspired Representation. J. Phys. Chem. Lett. 10, 4962-4967 (2019).
Zhang, Y., Ye, S., Zhang, J., Jiang, J. & Jiang, B. J. Phys. Chem. B 124, 7284-7290 (2020).
Zhang, Y., Xia, J. & Jiang, B. REANN: A PyTorch-based End-to-End Multi-functional Deep Neural Network Package. Phys. Rev. Lett. 127, 156002 (2021).
Zhang, Y., Xia, J. & Jiang, B. J. Chem. Phys. 156, 114801 (2022).

上游代码仓库：https://github.com/zhangylch/REANN
昇腾 NPU 适配版：https://atomgit.com/AI4Science/REANN

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

更多推荐

鲲鹏轻量级虚拟化实验（华为云ECS版）

鲲鹏昇腾开发者社区

华为昇腾NPU开发实战：YOLOv11从GPU到Atlas 300I Duo全流程部署

鲲鹏昇腾开发者社区

【Ascend】基于四台Atlas800T A2混合部署GLM5.3-w8a8c8

基于v0.30.0rc1镜像，采用混合部署方式，在多节点上通过Docker启动Ascend算力容器，挂载模型与驱动资源。主节点配置start_model_server.sh启动vLLM服务，设置网络接口、IP及分布式参数，实现4路数据并行、8路张量并行，启用前缀缓存、推测生成与FlashComm优化，支持大模型高效推理。