AtomicConv 模型迁移适配指导

1. 模型概述

AtomConvModel 是原子卷积神经网络，用于预测蛋白-配体结合亲和力。该模型将分子建模为 3D 空间中的近邻图，通过原子卷积层提取局部几何特征。模型接收 12 个输入张量分别对应两个分子片段（frag1、frag2）和完整复合物的坐标、邻居索引和原子序数，经过三个并行的 AtomicConvolution 层计算径向对称函数以捕获不同距离范围内的原子相互作用，随后将特征拼接后通过全连接层进行回归预测。该模型可灵活调整分子大小、卷积层参数和神经网络结构以适应不同任务需求。该模型在药物发现、材料科学和计算化学等任务中展现出强大能力，本文描述的 DMPNN Model 模型是基于 DeepChem 套件实现的，后续适配也是基于该套件修改。

2. 准备运行环境

2.1 软件环境

组件	版本
Python	3.10.19
PyTorch	2.1.0
torch_npu	2.1.0.post13
CANN	8.1.RC1

2.2 硬件环境

设备型号	NPU 配置
Atlas 800T A2	单卡 / 多卡

2.3 准备镜像

镜像环境	镜像地址
公网	swr.cn-southwest-2.myhuaweicloud.com/atelier/pytorch_2_1_ascend:pytorch_2.1.0-cann_8.1.rc1-py_3.10-euler_2.10.11-aarch64-snt9b-20250603154214-4e60e43

2.4 启动镜像

docker run -u root --privileged \
 --name {container_name} \
 --device /dev/davinci0 \
 --device /dev/davinci_manager \
 --device /dev/devmm_svm \
 --device /dev/hisi_hdc \
 -v /usr/local/dcmi:/usr/local/dcmi \
 -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
 -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
 -v /etc/ascend_install.info:/etc/ascend_install.info \
 -itd {image_id} /bin/bash

3 运行指导

3.1 安装已完成迁移的deepchem-ascend

docker exec -it {container_name} bash

直接安装 deepchem-ascend 二进制包，该包已基于适配代码重新编译并上传至 PyPI。

pip install deepchem-ascend==0.0.1

3.2 安装依赖

安装运行必要的依赖库

pip install torch-geometric

3.3 测试验证

测试代码举例如下：

import os
import numpy as np

from deepchem.data import NumpyDataset
from deepchem.models.torch_models.torch_model import is_npu_available
from deepchem.models.torch_models import AtomConvModel

try:
    import torch
    import torch_npu
except ModuleNotFoundError:
    print("torch_npu not installed, cannot test NPU")
    exit(0)

if not is_npu_available():
    print("NPU not available, skipping test")
    exit(0)

N_atoms = 5
batch_size = 1
frag1_num_atoms = N_atoms
frag2_num_atoms = N_atoms
complex_num_atoms = 2 * N_atoms

model = AtomConvModel(
    n_tasks=1,
    batch_size=batch_size,
    layer_sizes=[10],
    frag1_num_atoms=frag1_num_atoms,
    frag2_num_atoms=frag2_num_atoms,
    complex_num_atoms=complex_num_atoms,
    dropouts=0.0,
    learning_rate=0.003
)

features = []
frag1_coords = np.random.rand(N_atoms, 3)
frag1_nbr_list = {i: [] for i in range(N_atoms)}
frag1_z = np.random.randint(10, size=(N_atoms))
frag2_coords = np.random.rand(N_atoms, 3)
frag2_nbr_list = {i: [] for i in range(N_atoms)}
frag2_z = np.random.randint(10, size=(N_atoms))
system_coords = np.random.rand(complex_num_atoms, 3)
system_nbr_list = {i: [] for i in range(complex_num_atoms)}
system_z = np.random.randint(10, size=(complex_num_atoms))

features.append(
    (frag1_coords, frag1_nbr_list, frag1_z, frag2_coords, frag2_nbr_list,
     frag2_z, system_coords, system_nbr_list, system_z))
features = np.asarray(features, dtype=object)
labels = np.random.rand(batch_size)
train = NumpyDataset(features, labels)

model.fit(train, nb_epoch=10)

preds = model.predict(train)
print(f"Predictions shape: {preds.shape}")

model_device = next(model.model.parameters()).device
print(f"Model device: {model_device}")

assert model_device.type == 'npu', f"Model should be on NPU, but got {model_device}"
print("NPU test passed!")

测试结果：

1. 模型概述

组件

版本

Python

3.10.19

PyTorch

2.1.0

torch_npu

2.1.0.post13

CANN

8.1.RC1

设备型号

NPU 配置

Atlas 800T A2

单卡 / 多卡

镜像环境

镜像地址

公网

swr.cn-southwest-2.myhuaweicloud.com/atelier/pytorch_2_1_ascend:pytorch_2.1.0-cann_8.1.rc1-py_3.10-euler_2.10.11-aarch64-snt9b-20250603154214-4e60e43

2.4 启动镜像

docker run -u root --privileged \
 --name {container_name} \
 --device /dev/davinci0 \
 --device /dev/davinci_manager \
 --device /dev/devmm_svm \
 --device /dev/hisi_hdc \
 -v /usr/local/dcmi:/usr/local/dcmi \
 -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
 -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
 -v /etc/ascend_install.info:/etc/ascend_install.info \
 -itd {image_id} /bin/bash

3.3 测试验证

测试代码举例如下：

import os
import numpy as np

from deepchem.data import NumpyDataset
from deepchem.models.torch_models.torch_model import is_npu_available
from deepchem.models.torch_models import AtomConvModel

try:
    import torch
    import torch_npu
except ModuleNotFoundError:
    print("torch_npu not installed, cannot test NPU")
    exit(0)

if not is_npu_available():
    print("NPU not available, skipping test")
    exit(0)

N_atoms = 5
batch_size = 1
frag1_num_atoms = N_atoms
frag2_num_atoms = N_atoms
complex_num_atoms = 2 * N_atoms

model = AtomConvModel(
    n_tasks=1,
    batch_size=batch_size,
    layer_sizes=[10],
    frag1_num_atoms=frag1_num_atoms,
    frag2_num_atoms=frag2_num_atoms,
    complex_num_atoms=complex_num_atoms,
    dropouts=0.0,
    learning_rate=0.003
)

features = []
frag1_coords = np.random.rand(N_atoms, 3)
frag1_nbr_list = {i: [] for i in range(N_atoms)}
frag1_z = np.random.randint(10, size=(N_atoms))
frag2_coords = np.random.rand(N_atoms, 3)
frag2_nbr_list = {i: [] for i in range(N_atoms)}
frag2_z = np.random.randint(10, size=(N_atoms))
system_coords = np.random.rand(complex_num_atoms, 3)
system_nbr_list = {i: [] for i in range(complex_num_atoms)}
system_z = np.random.randint(10, size=(complex_num_atoms))

features.append(
    (frag1_coords, frag1_nbr_list, frag1_z, frag2_coords, frag2_nbr_list,
     frag2_z, system_coords, system_nbr_list, system_z))
features = np.asarray(features, dtype=object)
labels = np.random.rand(batch_size)
train = NumpyDataset(features, labels)

model.fit(train, nb_epoch=10)

preds = model.predict(train)
print(f"Predictions shape: {preds.shape}")

model_device = next(model.model.parameters()).device
print(f"Model device: {model_device}")

assert model_device.type == 'npu', f"Model should be on NPU, but got {model_device}"
print("NPU test passed!")

测试结果：