Ascend-SACT/reinvent4
模型介绍文件和版本Pull Requests讨论分析
下载使用量0

REINVENT 4 NPU 迁移指南

本文档指导用户将 REINVENT 4 分子设计工具从 PyTorch GPU 环境迁移到 Ascend NPU 环境，完成环境配置、源码准备、权重准备、模型迁移、测试验证及精度对齐。

目录

  1. 环境准备
  2. 源码与权重准备
  3. 数据准备
  4. 模型迁移
  5. 测试验证
  6. 推理验证
  7. 精度对齐

1. 环境准备

1.1 硬件与系统要求

要求规格
NPU 设备Ascend 910B3
CANN 版本8.5.0
Python 版本3.10
系统架构Linux aarch64

1.2 安装 Conda

# 下载并安装 Miniconda
bash Miniconda3-latest-Linux-aarch64.sh
source ~/.bashrc

1.3 创建 Python 环境

conda create --name reinvent4 python=3.10
conda activate reinvent4

1.4 设置镜像源（可选）

# 配置 Conda 华为云镜像
conda config --add channels https://repo.huaweicloud.com/ascend/repos/conda
conda config --show channels

# 配置 pip 华为云镜像
pip config set global.index-url https://repo.huaweicloud.com/repository/pypi/simple
pip config set global.trusted-host repo.huaweicloud.com
git config --global http.version HTTP/1.1

1.5 安装依赖

# 创建工作目录，自定义路径
mkdir -p ~/reinvent4 && cd ~/reinvent4

# 克隆 REINVENT 4 官方源码
git clone git@github.com:MolecularAI/REINVENT4.git
# 启动 REINVENT 4 安装（CPU 模式安装基础依赖）
python install.py cpu

# 查看 reinvent 帮助
reinvent --help

# 安装 rdkit
conda install -c conda-forge rdkit

# 补充 NPU 依赖
pip install decorator
pip install attrs
pip install psutil
pip install torch_npu

1.6 应用补丁

在项目源码目录中应用 NPU 适配补丁，修改 CUDA 相关代码为 NPU 设备调用：

cd ~/reinvent4/REINVENT4 #实际项目路径
git apply linkinvent_npu_patch_for_reinvent4.patch

1.7 替换 conda 安装路径源码文件

REINVENT 4 通过 python install.py cpu 安装后，实际运行时加载的是 conda 环境 site-packages 中的源码。因此，必须将补丁修改后的源码文件替换到 conda 安装路径中，否则修改不会生效。

# conda 默认安装路径（LinkInvent 模型目录）CONDA_LINKINVENT_PATH=~/miniconda3/envs/reinvent4/lib/python3.10/site-packages/reinvent/models/linkinvent/
# 源码补丁后的路径
SRC_LINKINVENT_PATH=~/reinvent4/REINVENT4/reinvent/models/linkinvent/
# 备份 conda 原始文件（可选但推荐）
cp -r $CONDA_LINKINVENT_PATH ${CONDA_LINKINVENT_PATH}_backup
# 将补丁修改后的源码文件复制到 conda 安装路径，覆盖原始文件
cp -f $SRC_LINKINVENT_PATH/*.py $CONDA_LINKINVENT_PATH/
# 验证替换结果ls -la $CONDA_LINKINVENT_PATH

1.8 环境变量配置

# 取消设备可见性限制
unset ASCEND_RT_VISIBLE_DEVICES

# 设置 CANN 环境变量（必须执行）
export ASCEND_TOOLKIT_HOME=/usr/local/Ascend/cann-8.5.0
source $ASCEND_TOOLKIT_HOME/bin/setenv.bash

# 可选：内存优化（libtcmalloc.so 需单独安装）
export TASK_QUEUE_ENABLE=2
export LD_PRELOAD=$LD_PRELOAD:/root/miniconda3/lib/libtcmalloc.so

1.9 验证 NPU 可用性

python3 -c "import torch; import torch_npu; print(torch.npu.is_available())"

2. 权重准备

2.1 下载预训练权重（Prior Models）

所有公开的 prior 模型可从 Zenodo 下载：

https://zenodo.org/records/15641297

下载模型权重后，放置到项目目录下的 prior 文件夹中：

cd ~/reinvent4/REINVENT4
mkdir -p prior
# 将下载的权重文件放入 prior/ 目录

2.3 目录结构

~/reinvent4/
└── REINVENT4/                      # REINVENT 4 官方源码
    ├── reinvent/                    # 核心代码
    │   ├── models/                  # 模型定义
    │   ├── rl/                      # 强化学习模块
    │   ├── tl/                      # 迁移学习模块
    │   └── scoring/                 # 评分组件
    ├── prior/                       # 预训练权重目录
    │   └── *.model                  # prior 模型文件
    ├── install.py                   # 安装脚本
    ├── linkinvent_rnn_eva.py        # 测试脚本
    └── README.md                    # 项目说明

3. 数据准备

测试用例 linkinvent_rnn_eva.py 内置了 LinkInvent RNN 模型评估流程，无需额外数据准备。

4. 模型迁移

4.1 迁移原理

REINVENT 4 源码中使用了 PyTorch 的 CUDA 设备相关接口，需要将设备指定从 cuda 修改为 npu，使模型能够在 Ascend NPU 上运行。

4.2 迁移要点

原始代码修改后代码
.cuda().to(device)
device='cuda'device='npu'
torch.cuda.is_available()hasattr(torch, "npu") and torch.npu.is_available()
torch.device("cuda")torch.device("npu")
torch.device("cuda:0")torch.device("npu:0")

4.3 迁移步骤

  1. 导入 torch_npu：在入口脚本中添加 import torch_npu，注册 NPU 后端。
  2. 设备参数传递：通过命令行参数 --device npu 指定 NPU 设备，确保设备信息正确传递到模型加载和推理流程中。
  3. 权重加载设备验证：运行测试时打印权重加载的设备信息，确认模型参数已正确加载到 NPU 上。
  4. 优化性能，规避npu部分组合算子报错

5. 测试验证

5.1 NPU 测试

cd ~/reinvent4/REINVENT4

# 运行 LinkInvent RNN 评估测试
python linkinvent_rnn_eva.py --device npu --prior /home/reinvent/REINVENT4/prior/linkinvent.prior --prior /home/reinvent/REINVENT4/prior/linkinvent.prior

5.2 参数说明

参数说明默认值
--device设备类型 (npu / cpu / cuda)cpu
  • 控制台打印加载权重的设备信息，预期为 npu
  • 运行过程无明显报错
  • 模型推理正常完成，输出分子生成/评估结果

5.4 验收标准

检查项预期结果
权重加载设备npu
运行状态无报错，正常完成
输出结果分子生成/评估结果正常输出

6. 推理验证

6.1 使用 Prior 模型生成linker

# 使用linkinvent rnn prior 模型进行分子采样（示例配置）
python linkinvent_rnn_eva.py --device npu --test sample --total-samples 10000 --batch-size 1000 --prior /home/reinvent/REINVENT4/prior/linkinvent.prior

7. 精度对齐

7.1 验证方案

由于 REINVENT 4 是生成式分子设计工具，精度对齐主要关注以下方面：

验证维度验证内容验收标准
模型加载权重加载设备加载到 NPU，无精度损失
分子生成生成分子合法性SMILES likehood与 CPU 一致
采样分布生成分子的分布一致性与 CPU 结果分布对齐

7.2 验证步骤

  1. 设备验证：运行 linkinvent_rnn_eva.py --device npu，确认权重加载设备为 npu，无报错。

  2. CPU 基线对比：使用相同权重在 CPU 模式下运行，记录生成分子结果。

  3. NPU 对比：使用相同权重在 NPU 模式下运行，记录生成分子结果。

  4. 结果对比：比较 CPU 与 NPU 的生成结果，确认：

    • SMILES 合法性一致
    • 生成分子分布对齐
    • 评分组件计算结果一致

7.3 精度验收标准

  • 模型权重在 NPU 上加载成功，设备信息显示为 npu
  • 测试用例 linkinvent_rnn_eva.py --device npu 运行无报错
  • 生成分子质量与 CPU 基线一致

附录

A. 常见问题

问题解决方案
import torch_npu 失败确认已安装 torch_npu 包，并正确执行 CANN 环境变量配置
NPU 设备不可用检查 source $ASCEND_TOOLKIT_HOME/bin/setenv.bash 是否执行
内存不足启用 TASK_QUEUE_ENABLE=2 和 LD_PRELOAD 内存优化
openeye 安装失败确认使用 aarch64 架构的 whl 包，或使用 conda 安装
rdkit 未安装执行 conda install -c conda-forge rdkit

B. 环境信息汇总

组件版本
Python3.10
CANN8.5.0
PyTorch>= 2.0.0 (NPU 适配版)
torch_npu对应 CANN 版本
rdkitconda-forge 最新版
openeye2025.2.3
REINVENT4