CHGNet模型在昇腾平台的迁移适配实践

weixin_45870772

33人浏览 · 2026-07-30 10:35:04

weixin_45870772 · 2026-07-30 10:35:04 发布

作者：昇腾实战派

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

背景概述

在材料科学领域，密度泛函理论（DFT）是计算材料性质的金标准，但其高昂的计算成本严重制约了大规模材料筛选和分子动力学模拟的效率。机器学习势能（Machine Learning Potential, MLP）为这一问题提供了新的解决思路——通过在 DFT 数据上训练神经网络，实现接近 DFT 精度但快几个数量级的原子模拟。

CHGNet（Crystal Hamiltonian Graph neural Network）由加州大学伯克利分校 Ceder 课题组于 2023 年发表于 Nature Machine Intelligence，是一款电荷感知的通用神经网络势能模型。CHGNet 在 Materials Project 超过 150 万个结构（来自 14.6 万种化合物）的 GGA/GGA+U 静态和弛豫轨迹上完成预训练，覆盖整个元素周期表，能够同时预测能量、力、应力和磁矩，是目前材料稳定性预测领域的顶级模型之一。

CHGNet 的核心创新在于将电荷信息引入图神经网络——通过用 DFT 磁矩正则化原子特征，捕获局部离子环境和电荷分布的丰富信息，使模型能够准确描述电子相互作用，这对于过渡金属氧化物等电荷有序系统尤为关键。

本文介绍 CHGNet 的昇腾 Ascend NPU 适配版本——在保持原始模型精度的前提下，实现在昇腾 AI 平台上的结构预测和几何优化推理。

模型介绍

CHGNet 概述

CHGNet 是一个预训练的通用神经网络势能模型，输入晶体结构，输出：

能量（eV/atom）：体系总能量
力（eV/Å）：每个原子的受力
应力（GPa）：晶胞应力张量
磁矩（μ_B）：每个原子的磁矩

支持的应用场景：

结构弛豫（几何优化）
分子动力学（MD）模拟
声子频谱计算
弹性张量计算
材料稳定性筛选

整体架构

CHGNet 基于图神经网络架构，将晶体结构表示为两层图：原子图（atom graph）和键图（bond graph），分别捕获原子间相互作用和键角信息。

图构建

图类型	节点	边	截断半径	作用
原子图（Atom Graph）	原子	原子间键	6 Å	描述原子对相互作用
键图（Bond Graph）	键	共享原子的键对	3 Å	描述键角三体相互作用

模型组成

嵌入层
- AtomEmbedding：将原子序数映射为特征向量（64维）
- BondEncoder：使用径向基函数（RBF, 31个基）编码键长
- AngleEncoder：使用傅里叶基函数（31个基）编码键角
交互层（Interaction Blocks）×4

每个 Block 包含：
- AtomConv：原子卷积层，通过消息传递聚合邻居信息更新原子特征
- BondConv：键卷积层，通过键图消息传递更新键特征
- AngleUpdate：角度更新层，更新角度嵌入
使用 Gated MLP 作为非线性变换，SiLU 激活函数，支持 Layer Normalization。
读出层（Readout）
- GraphPooling：全局平均池化（或多头注意力池化）聚合原子特征为体系级表示
- MLP：多层感知机输出能量预测
- 力和应力通过能量对坐标的自动微分获得

电荷感知机制

CHGNet 的核心创新：用 DFT 计算的磁矩作为辅助训练目标正则化原子特征。磁矩携带了丰富的局部电荷分布信息——不同氧化态的过渡金属原子具有不同的磁矩，模型通过学习预测磁矩，间接学习了电荷分布和电子相互作用。

平滑截断函数

使用可学习的余弦截断函数确保能量对距离的导数（即力）在截断半径处连续，避免力场中的数值不连续。

预训练数据

来源：Materials Project 的 GGA/GGA+U 弛豫轨迹
规模：超过 150 万个结构，14.6 万种化合物
覆盖：整个元素周期表
标签：能量、力、应力、磁矩

可用预训练模型

版本	说明
`0.3.0`（默认）	MPtrj 预训练 CHGNet
`0.2.0`	旧版 MPtrj，向后兼容 NMI 论文
`r2scan`	R2SCAN 级别，从 MP-R2SCAN 数据集迁移学习

与 ASE 集成

CHGNet 提供 CHGNetCalculator，完全兼容 ASE（Atomic Simulation Environment），支持：

BFGS / FIRE / LBFGS 结构优化
VelocityVerlet / NPT 分子动力学
UnitCellFilter 晶胞优化

昇腾 NPU 适配版本

迁移动机

原始 CHGNet 基于 PyTorch 实现，默认支持 CPU 和 CUDA。本项目通过安装 torch_npu 并指定 use_device='npu'，即可在昇腾 Ascend NPU 上运行推理，无需修改模型代码。

仓库结构

chgnet_npu/
├── README.md
├── setup.py
├── pyproject.toml
├── npu_test.py                         # NPU 推理测试脚本
├── chgnet/
│   ├── model/
│   │   ├── model.py                    # CHGNet 主模型
│   │   ├── dynamics.py                 # CHGNetCalculator + MD 模拟
│   │   ├── layers.py                   # AtomConv/BondConv/AngleUpdate
│   │   ├── encoders.py                 # Atom/Bond/Angle 编码器
│   │   ├── basis.py                    # RBF/Fourier 基函数
│   │   ├── functions.py                # MLP/GatedMLP
│   │   └── composition_model.py        # 原子参考能量
│   ├── graph/
│   │   ├── converter.py                # 结构→图转换器
│   │   ├── crystalgraph.py             # CrystalGraph 数据结构
│   │   └── fast_converter_libraries/   # C 加速图构建
│   ├── data/
│   │   └── dataset.py                  # 数据集加载
│   ├── trainer/
│   │   └── trainer.py                  # 训练器
│   ├── pretrained/
│   │   ├── 0.2.0/                      # v0.2.0 权重
│   │   ├── 0.3.0/                      # v0.3.0 权重（默认）
│   │   └── r2scan/                     # R2SCAN 权重
│   └── utils/
├── examples/
│   ├── basics.ipynb                    # 基础用法
│   ├── fine_tuning.ipynb               # 微调示例
│   ├── run_md.py                       # MD 模拟
│   └── mp-18767-LiMnO2.cif            # 示例结构
└── tests/

适配方式

CHGNet 原始代码中的 determine_device() 函数已支持自动设备检测。NPU 适配只需：

安装 torch_npu
使用 CHGNetCalculator(use_device='npu') 指定设备

无需修改模型内部代码，torch_npu 自动将 PyTorch 算子调度到 NPU。

版本信息

软件	版本
HDK	25.2.0
CANN	8.3.RC1
Python	3.11
PyTorch	2.7.1
torch_npu	2.7.1

环境配置

创建 Conda 环境

conda create --name chgnet python=3.11
conda activate chgnet

克隆代码

git clone https://atomgit.com/AI4Science/chgnet.git

安装依赖

pip install chgnet
pip install pyyaml torch_npu==2.7.1 decorator attrs psutil absl-py cloudpickle ml-dtypes scipy tornado

主要依赖：

chgnet：CHGNet 模型包（含预训练权重）
torch_npu==2.7.1：昇腾 NPU 支持
pymatgen：材料学 Python 库（结构表示、CIF 解析）
ase：原子模拟环境（结构优化、MD）

模型权重

预训练权重已内置在 chgnet 包中，无需额外下载：

chgnet/pretrained/
├── 0.2.0/chgnet_0.2.0_e30f77s348m32.pth.tar
├── 0.3.0/chgnet_0.3.0_e29f68s314m37.pth.tar
└── r2scan/chgnet_r2scan_transfer_learning_e15f36s161m23.pth.tar

加载方式：

from chgnet.model.model import CHGNet

chgnet = CHGNet.load()              # 默认 v0.3.0
chgnet = CHGNet.load(model_name='r2scan')  # R2SCAN 版本

推理命令

NPU 结构弛豫

python npu_test.py

脚本内容：

from ase.io import read
from ase.optimize import BFGS
from ase.filters import UnitCellFilter
from chgnet.model import CHGNetCalculator

crystal = read("mp-18767-LiMnO2.cif")
crystal.calc = CHGNetCalculator(use_device='npu')
sf = UnitCellFilter(crystal)
optimizer = BFGS(sf)
optimizer.run(fmax=0.05, steps=100)

直接预测（静态计算）

from chgnet.model.model import CHGNet
from pymatgen.core import Structure

chgnet = CHGNet.load()
structure = Structure.from_file('examples/mp-18767-LiMnO2.cif')
prediction = chgnet.predict_structure(structure)

print(f"Energy: {prediction['e']:.4f} eV/atom")
print(f"Forces shape: {prediction['f'].shape}")
print(f"Stress shape: {prediction['s'].shape}")
print(f"Magmom shape: {prediction['m'].shape}")

分子动力学模拟

from chgnet.model.dynamics import MolecularDynamics

md = MolecularDynamics(
    atoms=structure,
    model=CHGNet.load(),
    ensemble="nvt",
    temperature=1000,  # K
    timestep=2,        # fs
    use_device='npu',
)
md.run(steps=1000)

模型特点与优势

与其他通用势能模型对比

模型	电荷感知	预训练数据规模	输出
CHGNet	✅ 磁矩正则化	150万结构/14.6万化合物	能量+力+应力+磁矩
M3GNet	❌	18.7万结构	能量+力+应力
MACE-MP-0	❌	15万结构	能量+力+应力

CHGNet 的电荷感知优势

能区分同一元素不同氧化态（如 Fe²⁺ vs Fe³⁺）
准确描述电荷有序/电荷转移驱动的相变
对过渡金属氧化物、电池材料等电荷敏感体系更准确
在 Matbench Discovery 材料稳定性预测基准上达到顶级性能

迁移适配要点

设备指定

# CPU 推理
calc = CHGNetCalculator(use_device='cpu')

# NPU 推理
calc = CHGNetCalculator(use_device='npu')

无侵入式适配

CHGNet 的 NPU 适配是零代码修改的——模型内部使用标准 PyTorch 算子，torch_npu 通过设备调度机制自动将计算迁移到 NPU，无需重写图神经网络消息传递逻辑。

注意事项

图构建（CrystalGraphConverter）在 CPU 上执行（涉及 pymatgen 结构操作）
模型推理（前向传播 + 自动微分求力）在 NPU 上执行
对于 MD 模拟，每步的图构建开销可能成为瓶颈

已知限制

图构建步骤仍在 CPU 上运行，对于大量结构的批量预测存在 CPU-NPU 数据传输开销
预训练数据基于 GGA/GGA+U，对于需要更高精度的场景建议使用 r2scan 版本
分子动力学模拟中每步需重建图（如原子位移超出截断范围），这部分暂无法加速

参考文献

Deng, B., Zhong, P., Jun, K., Riebesell, J., Han, K., Bartel, C.J. & Ceder, G. CHGNet as a pretrained universal neural network potential for charge-informed atomistic modelling. Nature Machine Intelligence 5, 1031-1041 (2023). https://nature.com/articles/s42256-023-00716-3
上游代码仓库：https://github.com/CederGroupHub/chgnet
昇腾 NPU 适配版：https://atomgit.com/AI4Science/chgnet
API 文档：https://chgnet.lbl.gov