MAT模型迁移适配指导

1. 模型概述

MATModel（Molecular Attention Transformer）是一种基于Transformer架构的分子性质预测模型，专为小分子量子化学任务（如FreeSolv数据集的水合自由能回归）设计。它将每个分子表示为三个核心张量：原子节点特征矩阵、邻接矩阵和原子间距离矩阵，并结合注意力机制显式建模原子间的结构与空间关系。

该模型需配合MATFeaturizer使用，后者将SMILES转换为包含3D几何信息的MATEncoding。输入经MATEmbedding、多层MATEncoder（含自注意力与距离感知注意力）和MATGenerator处理后，输出每个分子的标量预测值（如溶解自由能）。

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 启动镜像

IMAGE_ID=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
CONTAINER_NAME=matmodel
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 创建环境

docker exec -it ${CONTAINER_NAME} bash
conda create -n matmodel --clone PyTorch-2.1.0
conda activate matmodel

3.2 安装依赖

pip install 'numpy>=1.19.2,<=1.24.0' scipy pandas scikit-learn rdkit sympy attrs pathlib2 psutil pyyaml protobuf decorator requests absl-py tqdm pytest flaky 'transformers>=4.28,<4.35'

3.3 迁移适配

直接安装deepchem-ascend二进制包，已基于适配代码重新编译成二进制包上传pypi

pip install deepchem-ascend==0.0.1

3.4 示例

import numpy as np
import deepchem as dc

# load datasets
task, datasets, trans = dc.molnet.load_freesolv()
train, valid, test = datasets

# initialize model
model = dc.models.torch_models.MATModel(n_encoders=2,
                                        sa_hsize=128,
                                        d_input=128,
                                        d_hidden=128,
                                        d_output=128,
                                        encoder_hsize=128,
                                        embed_input_hsize=36,
                                        gen_attn_hidden=32)
# overfit test
model.fit(valid, nb_epoch=100)
metric = dc.metrics.Metric(dc.metrics.mean_absolute_error,
                          mode="regression")
scores = model.evaluate(valid, [metric], trans)
print(f"模型所在设备: {model.device}")
print(f"mean_absolute_error:{scores['mean_absolute_error']}")

3.5 运行测试代码

复制上述测试代码保存到test_mat.py

python test_mat.py

测试结果：

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 启动镜像

IMAGE_ID=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
CONTAINER_NAME=matmodel
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.4 示例

import numpy as np
import deepchem as dc

# load datasets
task, datasets, trans = dc.molnet.load_freesolv()
train, valid, test = datasets

# initialize model
model = dc.models.torch_models.MATModel(n_encoders=2,
                                        sa_hsize=128,
                                        d_input=128,
                                        d_hidden=128,
                                        d_output=128,
                                        encoder_hsize=128,
                                        embed_input_hsize=36,
                                        gen_attn_hidden=32)
# overfit test
model.fit(valid, nb_epoch=100)
metric = dc.metrics.Metric(dc.metrics.mean_absolute_error,
                          mode="regression")
scores = model.evaluate(valid, [metric], trans)
print(f"模型所在设备: {model.device}")
print(f"mean_absolute_error:{scores['mean_absolute_error']}")