作者：昇腾实战派

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

背景概述

量子化学计算是理解分子性质和化学反应的基础，但传统第一性原理方法（如密度泛函理论）的高昂计算成本制约了大规模分子体系的研究。深度学习为量子化学提供了全新的求解范式——通过在量子化学数据上训练神经网络，以极低的计算代价实现接近量子化学精度的分子性质预测。

SchNet 由柏林工业大学 Müller 课题组于 2017 年发表于 NeurIPS，是连续滤波卷积神经网络（continuous-filter convolutional neural network）的开创性工作，用于建模原子间量子相互作用。SchNet 的核心创新在于引入了连续滤波器——使用径向基函数将原子间距离扩展为连续的滤波器权重，使得模型对原子位置的变化具有连续、可微的响应，这从根本上不同于传统离散卷积。

SchNet 的数学形式为：

$${\mathbf{x}}_i^{\prime }=\sum _{j\in \mathcal{N}(i)}{\mathbf{x}}_j\odot h_{\Theta }(\exp (-\gamma ({\mathbf{e}}_{j,i}-\mathbf{\mu })))$$

其中 $h_{\Theta }$ 为 MLP，${\mathbf{e}}_{j,i}$ 为原子间距离，$\mathbf{\mu }$ 为高斯中心。

本文介绍 SchNet 的昇腾 Ascend NPU 适配版本——通过自定义 AscendC 算子实现 radius_graph 高性能计算，并在 QM9 数据集上完成端到端训练。

模型介绍

SchNet 概述

SchNet 是一个端到端的分子性质预测模型，输入原子类型和三维坐标，输出分子量子化学性质（如能量、偶极矩、HOMO/LUMO 等）。模型的关键特点：

• 旋转不变性：仅使用原子间距离（标量），不依赖坐标系选择
• 连续可微：对原子位置连续可微，可通过自动微分获得原子力
• 局域性：通过截断半径限制相互作用范围
• 通用性：适用于分子和周期性体系

整体架构

模型组成

1. 原子嵌入层（Embedding）

   • 将原子序数 Z 映射为 128 维特征向量
   • 支持 padding（Z=0）用于批处理

2. 径向交互图（RadiusInteractionGraph）

   • 基于截断半径（默认 10 Å）构建原子对图
   • 限制最大邻居数（默认 32）
   • 输出边索引和边权重（原子间距离）

3. 高斯展开（GaussianSmearing）

   • 将标量距离展开为 50 个高斯基函数
   • 均匀分布在 [0, cutoff] 范围内
   • 为连续滤波器提供多尺度距离表示

4. 交互层（InteractionBlock）×6

   每个 Block 包含：

   • 连续滤波卷积（CFConv）：
     • 滤波器网络：Distance → MLP → Filter weights
     • 余弦截断函数：$C=0.5(\cos (\pi r/r_c)+1)$ 确保平滑截断
     • 消息传递：$m_{j\to i}=x_j\odot W_{ij}$
   • ShiftedSoftplus 激活函数：$f(x)=\text{softplus}(x)-\ln 2$
   • 线性层：更新原子特征

5. 读出层

   • 两层 MLP（128 → 64 → 1）
   • 全局聚合（sum 或 mean pooling）
   • 可选原子参考能量修正（atomref）

连续滤波卷积（CFConv）

这是 SchNet 的核心创新。传统图卷积使用离散的边类型或固定权重，而 CFConv 通过 MLP 将连续距离映射为滤波器权重：

距离 r_ij → 高斯展开 → MLP → 滤波器权重 W_ij
消息 = x_j * W_ij * Cutoff(r_ij)
聚合 = Σ 消息

这使得模型能够学习任意形状的径向相互作用函数，而非预定义的固定函数形式。

QM9 数据集

QM9 是量子化学领域的标准基准数据集：

• 约 13 万个小有机分子（≤9 个重原子：C, H, O, N, F）
• 12 个量子化学性质（能量、偶极矩、电子性质等）
• DFT(B3LYP/6-31G(2df,p)) 计算

应用场景

• 分子性质预测：能量、力、偶极矩等
• 分子动力学：通过能量对坐标求导获得原子力
• 势能面扫描：高效探索分子构象空间
• 高通量筛选：快速评估候选分子性质

昇腾 NPU 适配版本

迁移动机

SchNet 的核心计算瓶颈在于 radius_graph——基于截断半径搜索原子邻居对。原始实现使用 torch_cluster.radius 在 CPU 上执行邻居搜索，这是训练循环中的性能瓶颈。本项目通过自定义 AscendC 算子将 radius 操作实现为 NPU 原生算子，实现全流程 NPU 加速。

仓库结构

SchNet_npu/
├── README.md
├── QM9.py                              # 主训练脚本（SchNet + QM9）
├── install.sh                          # 算子安装脚本
├── Radius.json                         # Radius 算子定义（输入输出规格）
├── Radius/                             # ⭐ AscendC Radius 算子实现
│   ├── op_kernel/                      # 算子核函数（AscendC 代码）
│   └── op_host/                        # 算子 Host 端逻辑
├── CustomOp/                           # 编译产物：自定义算子包
│   ├── build.sh
│   ├── op_kernel/
│   ├── op_host/
│   ├── framework/
│   └── build_out/                      # .run 安装包
└── CppExtensionInvocation/             # PyTorch C++ 扩展封装
    ├── setup.py                        # 构建 custom_ops_lib
    ├── csrc/                           # C++ binding 源码
    └── custom_ops/                     # Python 接口

核心迁移改动

1. 自定义 AscendC Radius 算子

radius_graph 是 SchNet 最关键的操作——对每个原子，搜索截断半径内的所有邻居原子。原始 torch_cluster.radius 仅支持 CPU/CUDA，本项目实现了 NPU 原生的 Radius 算子：

# 原始实现（CPU）：
edge_index = torch_cluster.radius_graph(pos, r=cutoff, batch=batch)

# NPU 实现：
edge_index = torch.ops.myops.my_op(x, x, ptr, ptr, r, max_num_neighbors, False)

算子特性：

• 使用 AscendC 编程模型，充分利用 NPU AI Core 计算能力
• 支持批处理（通过 ptr 指定每个分子的原子范围）
• 支持 max_num_neighbors 限制
• 输入输出为 float32 ND 张量

2. 算子安装流程

# 生成算子框架
msopgen gen -i Radius.json -c ai_core-Ascend910B3 -lan cpp -out CustomOp

# 编译算子
cd CustomOp && bash build.sh

# 安装算子包
cd build_out && ./custom_opp_openEuler_aarch64.run

# 构建 PyTorch C++ 扩展
cd CppExtensionInvocation && bash build_and_run.sh

3. NPU 训练适配

import torch_npu
import custom_ops_lib  # 加载自定义算子

device = torch.device('npu')
model = model.to(device)

版本信息

软件	版本
CANN	8.1.RC1
HDK	25.2.0
Python	3.10
PyTorch	2.1.0
torch_npu	2.1.0.post17
scikit-learn	1.1.3

环境配置

创建容器

docker run -it -u root \
--net=host --shm-size=5g \
--device=/dev/davinci_manager \
--device=/dev/hisi_hdc \
--device=/dev/davinci0 \
--device=/dev/davinci1 \
-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 schnet_test \
swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.1.rc1-910b-openeuler24.03-py3.10

安装 PyTorch + torch_npu

wget https://download.pytorch.org/whl/cpu/torch-2.1.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
pip3 install torch-2.1.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
pip3 install torch_npu-2.1.0.post17-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

安装系统依赖

yum install -y gcc gcc-c++ libstdc++-devel libstdc++ make cmake openblas-devel openblas-static util-linux

安装 PyG 组件

pip3 install torch_cluster torch_geometric --no-build-isolation
pip3 install ase

安装 SchNet 自定义算子

git clone https://atomgit.com/AI4Science/SchNet.git
cd SchNet
export CPLUS_INCLUDE_PATH=/usr/include/c++/12/:$CPLUS_INCLUDE_PATH
bash install.sh -v Ascend910B3
cd CustomOp/build_out
./custom_opp_openEuler_aarch64.run
cd ../../CppExtensionInvocation/
bash build_and_run.sh

训练

source /usr/local/Ascend/ascend-toolkit/set_env.sh
python3 QM9.py

训练配置：

• 数据集：QM9（~130k 分子）
• 划分：50k 测试 + 1k 验证 + 其余训练
• Batch size：32
• 优化器：Adam，lr=0.001
• 学习率调度：ReduceLROnPlateau（patience=5, factor=0.7）
• 模型参数：hidden=128, filters=128, interactions=6, gaussians=50, cutoff=10.0

输出示例：

Epoch: 001, LR: 0.001000, Loss: 0.XXXXXXX, Val MAE: X.XXXXX ± X.XXXXX, Test MAE: X.XXXXX ± X.XXXXX

迁移适配要点

自定义算子开发流程

步骤	说明
1. 定义算子规格	Radius.json 描述输入输出 tensor 格式和类型
2. 生成框架	msopgen 自动生成 AscendC 算子工程骨架
3. 实现核函数	在 op_kernel/ 中编写 AscendC 并行计算逻辑
4. 编译安装	build.sh → .run 安装包 → 注册到 CANN
5. PyTorch 绑定	C++ Extension 将算子暴露为 torch.ops.myops.my_op

关键技术点

• Radius 算子：基于截断半径的邻居搜索，输出边索引对。AscendC 实现充分利用 NPU 并行计算能力，相比 CPU 串行搜索大幅加速
• Batch 支持：通过 ptr（batch 指针）区分不同分子的原子，避免跨分子错误连边
• PyTorch Geometric 兼容：保持与 PyG MessagePassing 框架的接口兼容，CFConv 继承自 MessagePassing

与原始实现的差异

组件	原始实现	NPU 实现
radius_graph	torch_cluster.radius（CPU/CUDA）	torch.ops.myops.my_op（NPU AscendC）
设备	CUDA GPU	Ascend NPU
消息传递	PyG MessagePassing	同（scatter_add 等已由 torch_npu 支持）

已知限制

• 自定义算子需要针对具体 SoC 型号编译（如 Ascend910B3）
• 当前 radius 算子固定使用 float32，不支持 float16
• QM9 数据集下载需要网络连接
• 大分子体系（原子数 > 1000）的邻居搜索可能需要调整 max_num_neighbors

参考文献

• Schütt, K.T., Kindermans, P.-J., Sauceda, H.E., Chmiela, S., Tkatchenko, A., Müller, K.-R. SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. NeurIPS 30, pp. 992-1002 (2017).
• Schütt, K.T., Arbabzadah, F., Chmiela, S., Müller, K.-R., Tkatchenko, A. Quantum-chemical insights from deep tensor neural networks. Nature Communications 8, 13890 (2017). doi: 10.1038/ncomms13890
• 上游代码仓库：https://github.com/atomistic-machine-learning/SchNet
• PyTorch Geometric SchNet 实现：https://pytorch-geometric.readthedocs.io/
• 昇腾 NPU 适配版：https://atomgit.com/AI4Science/SchNet