chemprop模型在昇腾上的迁移部署实践

作者：昇腾实战派

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

背景概述

分子性质预测是药物发现与材料设计中的高频环节：候选化合物的溶解性、毒性、PCE 等性质的实验测定成本高昂，DFT 等量子化学计算又受限于体系规模。深度学习为此提供了新的求解范式——将分子表示为图结构，用图神经网络端到端地学习"分子结构 → 性质"的映射，在海量历史数据上训练后即可对新分子进行快速性质评估。

chemprop 由 MIT 课题组开发（Yang et al., 2019 发表于 J. Chem. Inf. Model.），采用 MIT 协议开源，是分子性质预测领域应用最广泛的开源工具之一，其核心是有向消息传递神经网络（D-MPNN，directed Message Passing Neural Network）。团队于 2023 年发表的专文 Chemprop: A Machine Learning Package for Chemical Property Prediction（J. Chem. Inf. Model.）对其功能与基准测试做了系统描述。该工具在真实药物发现中已产出标志性成果：

• Halicin 的发现（Cell, 2020）：用 chemprop 预测大肠杆菌抗生素活性，从筛选库中发现全新结构类别的抗生素候选 Halicin
• 可解释的抗 MRSA 抗生素发现（Nature, 2023）：利用 chemprop 模型集成定位对耐药菌有效的新型结构类别，并用其 interpret 功能进行解释
• ADMET-AI 平台：在 Therapeutics Data Commons 的 41 个 ADMET 数据集上训练 chemprop 模型，提供大规模化合物库的吸收/分布/代谢/排泄/毒性评估服务

其 D-MPNN 与普通 MPNN 在原子节点上传递消息不同，是在有向边上传递消息，并在聚合时排除反向边，避免"信息回流"导致的过平滑：

$$h_{wv}^0=x_w$$

$$h_{wv}^{t+1}=\text{ReLU}\left(h_w^0+\sum _{k\in {\mathcal{N}}_w\setminus v}{h}_{kw}^t\right)$$

$$h_v=\text{ReLU}\left(W_o\sum _{w\in {\mathcal{N}}_v}{h}_{wv}^T\right)$$

其中 $h_{wv}$ 为沿边 $w\to v$ 的消息，$x_w$ 为原子特征。

本文介绍 chemprop 的昇腾 Ascend NPU 适配版本——通过 torch_npu 提供的 transfer_to_npu 兼容层以最小侵入方式完成迁移，仅修改 2 行代码，即可在 MoleculeNet 分类/回归数据集上完成端到端 NPU 训练。

模型介绍

chemprop 概述

chemprop 是面向分子性质预测的完整命令行工具集（本适配基于 chemprop v1.6.1），主要特点：

• 任务类型齐全：支持 classification、regression、multiclass、spectra 四类任务，天然支持多任务学习（缺失目标值自动掩码，不参与损失与评估）
• 损失与指标按任务适配：regression 默认 MSE 损失/RMSE 指标，classification 默认 BCE 损失/AUC 指标，multiclass 默认交叉熵，spectra 默认 SID 谱信息散度，均可通过参数切换（如 mcc、evidential、wasserstein 等）
• D-MPNN 编码器：有向边消息传递，配合 RDKit 原子/键特征
• 分子描述子特征：可拼接 rdkit_2d_normalized（200 维 CDF 归一化 RDKit 描述子）、morgan 指纹等特征，支持原子/键级自定义描述子
• 数据划分灵活：支持 random（默认 80/10/10）、scaffold_balanced、k 折交叉验证（–num_folds）、模型集成（–ensemble_size）等方式
• 不确定度估计：预测阶段支持 ensemble 方差、Monte Carlo dropout、mve、evidential 等方法及后续校准
• 工具链完整：训练（train.py）、预测（predict.py）、模型解释（interpret.py）、超参优化（hyperparameter_optimization.py）、传统 ML 基线（sklearn_train.py）、Web 界面（chemprop_web）

整体架构

模型组成

1. 数据处理（chemprop/data）

   • SMILES → RDKit 分子图 → 原子/键特征向量（one-hot + 数值特征拼接）
   • MoleculeDataset / MoleculeDataLoader 组织 batch（scaler 对目标与附加特征标准化）

2. D-MPNN 编码器（chemprop/models/mpn.py）

   • 在有向边维度上迭代更新消息（默认 depth=3）
   • 隐藏维度 hidden_size=300
   • 边消息聚合到节点，得到分子级表示

3. FFN 读出层（chemprop/models/ffn.py）

   • 两层 MLP（ffn_num_layers=2），activation + dropout
   • 节点表示聚合为分子向量：默认 mean，可选 sum 或 norm（默认 aggregation_norm=100）
   • 输出维度等于任务数（多任务时每任务一列）

4. 损失与优化（chemprop/train）

   • 按任务类型自动选择损失函数（BCE / MSE / 多分类交叉熵 / SID 光谱损失等）
   • Adam 优化器 + NoamLR 学习率调度（含 warmup，默认 warmup_epochs=2.0）

D-MPNN 与普通 MPNN 的区别

普通 MPNN 的消息更新在节点上进行，聚合来自所有邻居的信息；D-MPNN 将消息定义在有向边上，且在更新边 $w\to v$ 的消息时排除反向边 $v\to w$，使信息只能"单向流动"一轮再聚合到节点。这一设计避免了消息在无向图上来回穿梭，使分子表示学习更加稳定。

MoleculeNet 数据集

按官方说明，仓库 data.tar.gz 内置已预处理的 MoleculeNet 全套基准数据集与来自 LSC 的 ChEMBL 45 万分子子集（csv 格式）：

• 分类：hiv、sider、toxcast、muv 等（targets 为 0/1 标签，默认评价指标 AUC）
• 回归：delaney（水溶性）、qm8、qm9（量子化学能量性质）等（默认评价指标 RMSE/MAE）

每行一条样本：SMILES 列 + 若干目标性质列；多任务训练时缺失值留空即可，损失函数中自动掩码。

应用场景

• 药物发现：毒性、溶解性、生物活性等多性质端点预测（Halicin、抗 MRSA 抗生素等标志性案例）
• ADMET 评估：大规模化合物库的吸收/分布/代谢/排泄/毒性批量预测
• 性质筛选：虚拟化合物库的快速性质评估与优先级排序
• 光谱预测：NMR/UV-vis 等光谱任务（spectra 模式）
• 反应性质预测：经凝聚图反应（CGR）表示支持原子映射反应 SMILES 的性质预测
• 基线与研究：作为图神经网络分子表征的强基线，支持特征消融、超参搜索与不确定度研究

昇腾 NPU 适配版本

迁移动机

chemprop v1.6.1 的设备管理完全基于 torch.cuda API：

• chemprop/args.py 中 CommonArgs.device 属性返回 torch.device('cuda', self.gpu)，cuda 属性依赖 torch.cuda.is_available()
• 模型与数据通过 .to(device) 搬运，训练循环中无任何设备相关的分支代码

若逐文件将这些调用改写为 torch_npu API，改动点分散、侵入性强。torch_npu 提供的 torch_npu.contrib.transfer_to_npu 兼容层可在导入时一次性将 torch.cuda 系列 API（is_available、device、Tensor.cuda 等）映射为 NPU 对应实现，使上游代码零修改即可运行在 NPU 上。

本适配版本的目标是：全库仅增加 2 行 import，训练/预测链路自动驻留 NPU。

仓库结构

chemprop/
├── train.py                        # 训练入口（chemprop_train）
├── predict.py                      # 推理入口
├── interpret.py                    # 模型解释
├── hyperparameter_optimization.py  # 超参搜索
├── sklearn_train.py / sklearn_predict.py  # 传统 ML 基线
├── fingerprint.py                  # 分子指纹提取
├── chemprop/
│   ├── args.py                     # 参数定义（device/cuda 属性）
│   ├── data/                       # 数据读取、图构建、划分
│   ├── features/                   # 原子/键特征、分子描述子
│   ├── models/                     # MoleculeModel / D-MPNN / FFN
│   ├── train/
│   │   ├── train.py                # ⭐ NPU 适配改动点（2 行 import）
│   │   ├── run_training.py         # 单折训练流程
│   │   ├── cross_validate.py       # K 折交叉验证入口
│   │   ├── predict.py / evaluate.py / loss_functions.py
│   ├── utils.py / rdkit.py / spectra_utils.py
├── tests/data/                     # 测试数据（npz 格式）
├── data.tar.gz                     # MoleculeNet csv 数据集
└── setup.py

核心迁移改动

1. 引入 torch_npu 兼容层（全部代码改动）

改动仅位于 chemprop/train/train.py 头部：

# 原始实现：
import torch

# NPU 实现（增加 2 行）：
import torch_npu
from torch_npu.contrib import transfer_to_npu

选择该文件的原因：

• chemprop/__init__.py 导入 chemprop.train，后者导入 .train 模块——train.py 是所有入口（train/predict/interpret）共用导入链上的汇聚点，补丁在此加载后对整个进程生效
• transfer_to_npu 将 torch.cuda.is_available()、torch.device('cuda', gpu)、Tensor.cuda() 等调用透明映射到 NPU，因此：

# chemprop/args.py —— 无需修改
@property
def device(self) -> torch.device:
    if not self.cuda:                       # transfer_to_npu 使其在 NPU 主机上按 NPU 语义求值
        return torch.device('cpu')
    return torch.device('cuda', self.gpu)   # 自动映射为 NPU 设备

• 模型搬运（model.to(device)）、数据搬运（batch_batch_graph().to(device)）、优化器创建等代码路径全部保持原样，自动运行在 NPU 上

2. 依赖版本对齐

chemprop 对部分依赖有版本约束，需在 NPU 环境中显式对齐：

pip3 install astartes aimsim configargparse lightning>=2.0 \
    scikit-learn==1.1.3 descriptastorus rich cloudpickle ml-dtypes \
    tornado numpy==1.26.4

其中 numpy==1.26.4 与 torch 2.1.0 匹配，cloudpickle==1.6.1 为 checkpoint 兼容性所需。

版本信息

软件	版本
CANN	8.3.rc1
HDK	25.2.0
Python	3.11.13
PyTorch	2.1.0
torch_npu	2.1.0.post17
cloudpickle	1.6.1
scikit-learn	1.1.3
numpy	1.26.4

环境配置

创建容器

docker run -it -u root \
--net=host --shm-size=5g \
--device=/dev/davinci_manager \
--device=/dev/devmm_svm \
--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 \
-v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \
-v /home:/home/ \
--name chemprop_test \
--entrypoint=/bin/bash \
-it swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-910b-ubuntu22.04-py3.11

镜像获取可参考：昇腾镜像仓库-昇腾社区

安装 torch & torch_npu

wget https://download.pytorch.org/whl/cpu/torch-2.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
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-2.1.0-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

安装系统依赖库

apt update && apt install -y libsm6 libxext6

安装其它依赖

pip3 install astartes aimsim configargparse lightning>=2.0 scikit-learn==1.1.3 descriptastorus rich cloudpickle ml-dtypes tornado numpy==1.26.4

克隆代码仓

git clone https://atomgit.com/AI4Science/chemprop.git
cd chemprop && pip install -e .

训练

source /usr/local/Ascend/ascend-toolkit/set_env.sh
python3 train.py --data_path tests/data/classification.csv --dataset_type classification --save_dir classification_checkpoints

参数说明：

• <path>：csv 数据集路径
• <type>：[classification, regression, multiclass, spectra] 四种数据集类型之一
• <dir>：权重保存目录

数据说明：若使用分类数据集，可从仓库根目录 data.tar.gz 解压获取（data/ 下 hiv.csv、sider.csv、toxcast.csv、muv.csv 等），tests/data/ 内置的为 npz 格式测试数据。

训练配置：

• 任务类型：classification（默认评价指标 AUC；regression 默认 RMSE/MAE）
• 划分：默认 scaffold 划分 + 单折（cross_validate 驱动，可 --num_folds 扩展）
• Batch size：50
• Epochs：30（含 2.0 warmup epochs，NoamLR 调度）
• 模型参数：hidden_size=300，depth=3，ffn_num_layers=2

训练过程中 run_training 逐折打印各 epoch 损失与验证指标，训练完成后在 save_dir 下生成 model_0/model.pt 等 checkpoint 文件，并输出测试集指标（Metric: auc）。

迁移适配要点

迁移流程

步骤	说明
1. 依赖分析	定位设备 API 使用点（args.py 的 device/cuda 属性、模型与数据搬运）
2. 引入 torch_npu	安装与 torch 2.1.0 配套的 torch_npu-2.1.0.post17
3. 加载兼容层	train.py 头部增加 import torch_npu + transfer_to_npu，进程级生效
4. 依赖版本对齐	numpy1.26.4、scikit-learn1.1.3、cloudpickle==1.6.1 等
5. 容器化验证	CANN 镜像容器内完成 classification 任务端到端训练

关键技术点

• transfer_to_npu 机制：导入时对 torch API 做模块级替换，torch.cuda.is_available()、torch.device('cuda')、Tensor.cuda() 等自动指向 NPU 实现，适合这类设备调用分散、但统一走 torch.cuda 语义的存量代码
• 补丁加载位置的选择：改动点放在 chemprop/train/train.py（包导入链的公共节点），使 train.py、predict.py、interpret.py 等所有入口共享同一份补丁，无需各入口重复修改
• CPU 回退兼容：--no_cuda 参数依旧有效，可在无 NPU 环境中回退 CPU 运行
• 版本强绑定：torch 2.1.0 ↔ torch_npu 2.1.0.post17 ↔ CANN 8.3.rc1 三者需严格配套

与原始实现的差异

组件	原始实现	NPU 实现
import 层	仅 import torch	增加 import torch_npu + from torch_npu.contrib import transfer_to_npu
设备选择	torch.cuda.is_available() → CUDA GPU	同一调用经兼容层映射 → Ascend NPU
模型/数据搬运	.to(device)	同（无需修改）
训练命令	python3 train.py ...	完全一致
代码改动量	—	2 行 import，其余零修改

已知限制

• torch_npu 与 torch、CANN 版本强绑定，更换 torch 版本时需同步更换 torch_npu 与 CANN 镜像
• 补丁通过 chemprop 包的导入链加载；若绕开包导入直接以脚本方式引用单个模块，需自行导入 transfer_to_npu
• spectra、multiclass 等任务分支与 interpret/web 模块未在本环境做全量验证
• 仓库 tests/data/ 提供的是 npz 格式测试数据，README 命令中的 csv 数据集需从 data.tar.gz 或 MoleculeNet 官方渠道获取

参考文献

• Yang, K., Swanson, K., Jin, W., Coley, C., Eiden, P., Gao, H., Guzman-Perez, A., Hopper, T., Kelley, B., Mathea, M., Palmer, A., Settels, V., Jaakkola, T., Jensen, K., Barzilay, R. Analyzing Learned Molecular Representations for Property Prediction. J. Chem. Inf. Model. 59, 3370-3388 (2019).
• Yang, K., Swanson, K., Jin, W., Coley, C., Eiden, P., Gao, H., Guzman-Perez, A., Hopper, T., Kelley, B., Mathea, M., Palmer, A., Settels, V., Jaakkola, T., Jensen, K., Barzilay, R. Are Learned Molecular Representations Ready for Practice? A Perspective on Machine Learning-Based Scientific Software. arXiv:2011.01525.
• Heid, E., Greenman, K. P., Chung, Y., Li, S.-C., Graff, D. E., Vermeire, F. H., Wu, H., Green, W. H., McGill, C. J. Chemprop: A Machine Learning Package for Chemical Property Prediction. J. Chem. Inf. Model. 64, 9-17 (2024). doi: 10.1021/acs.jcim.3c01250