deepchem-ascend 0.0.8

Deep learning models for drug discovery, quantum chemistry, and the life sciences. The ascend + torch version.

1. 关于 deepchem-ascend

本项目基于开源deepchem，目标是让使用deepchem的模型能原生在昇腾上运行。当前受限于生态原因，deepchem社区暂时不接受对于ascend支持的代码修改合入，因此先在昇腾社区完成移植和独立版本发布，后续合入上游的deepchem代码仓。

所有开发工作将在develop分支以PR工作流形式运行，master分支仅用于和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 启动镜像

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=deepchem-ascend
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二进制包，已基于适配代码重新编译成二进制包上传pypi

pip install deepchem-ascend==0.0.8

3.2 下载模型源码

git clone --branch develop https://gitcode.com/AI4Science/deepchem-ascend.git

3.3 分布式训练示例脚本

将以下代码保存在test_train.py

import os
import deepchem as dc
from deepchem.models.lightning.trainer import LightningTorchModel
import torch
rank = int(os.getenv("RANK", "0"))
world_size = int(os.getenv("WORLD_SIZE", "1"))
local_rank = int(os.getenv("LOCAL_RANK", "0"))
print(f"当前进程 RANK: {rank}, 全局进程数: {world_size}, 本地设备号: {local_rank}")
tasks, datasets, _ = dc.molnet.load_clintox()
_, valid_dataset, _ = datasets
model = dc.models.MultitaskClassifier(
     n_tasks=len(tasks),
     n_features=1024,
     layer_sizes=[1000],
     dropouts=0.2,
     learning_rate=0.0001,
     device="npu",
     batch_size=16
)
trainer = LightningTorchModel(
             model=model,
             batch_size=16,
             devices=world_size,
             strategy="ddp_npu", #当前策略仅支持DDP
             accelerator="npu",  #加速器支持"npu"
             log_every_n_steps=1,
             fast_dev_run=False)
trainer.fit(valid_dataset, nb_epoch=30, num_workers=0) #当前仅支持num_workers=0

3.4 多卡训练执行命令

#nproc_per_node可以指定几卡训练，以下命令使用2张NPU
torchrun --nproc_per_node=2 --master_port=29500 test_train.py

4. 模型适配详情

当前模型适配策略以可运行为主要目标，精度和性能会在适配后按需进行调优。

模型	用途	框架适配	单卡	多卡	精度	性能	适配方法
Unet	用于图像分割的UNet模型，支持分类及回归任务	完成	已验证	未验证	未验证	未验证	UnetModel
ChemBERTa	Chemberta是一种基于SMILES字符串学习的Transformer风格模型。该模型架构基于RoBERTa架构设计，既可用于预训练生成嵌入表示，也可在下游应用中通过微调进行迁移学习。	完成	已验证	未验证	未验证	未验证	ChemBERTa
ChemCeption	利用卷积神经网络（CNN）的表示能力来预测分子性质。	完成	已验证	未验证	未验证	未验证	ChemCeption
TextCNNModel	基于SMILES字符串的卷积神经网络（CNN）用于分子性质预测。	完成	已验证	未验证	未验证	未验证	TextCNNModel
DTNNModel	基于SMILES字符串的卷积神经网络（CNN）用于分子性质预测。	完成	已验证	未验证	未验证	未验证	DTNN
MATModel	基于 Transformer 注意力机制的分子性质预测模型。	完成	已验证	未验证	未验证	未验证	MATModel
DMPNNModel	基于消息传递神经网络（MPNN）的分子性质预测模型。	完成	已验证	未验证	未验证	未验证	DMPNN
GNNModular	基于模块化、可组合的图神经网络（GNN）建模框架可灵活构建适用于分子性质预测模型。	完成	已验证	未验证	未验证	未验证	GNNModular
progressiveMultitaskModel	渐进式多任务学习模型，通过分阶段训练共享主干与任务特定子网络，有效缓解任务冲突，提升分子多性质预测性能。	完成	已验证	未验证	未验证	未验证	progressiveMultitask
Grover	一种面向分子图的自监督预训练模型。它将消息传递机制（Message Passing）与Transformer 架构深度融合，通过在大规模无标签分子数据上执行结构感知的预训练任务，学习通用且富含化学语义的分子表示。	完成	已验证	未验证	未验证	未验证	Grover
GANModel	一种由生成器和判别器组成的生成模型，两者通过对抗训练共同进化。	完成	已验证	未验证	未验证	未验证	GANModel
WGANModel	通过引入Wasserstein距离和权重裁剪机制改进的生成对抗网络。	完成	已验证	未验证	未验证	未验证	WGANModel
MolGANModel	一种结合了对抗训练和强化学习的图生成模型，能够直接生成有效的分子图结构，并可根据特定的化学性质目标进行定向优化。	完成	已验证	未验证	未验证	未验证	MolGANModel
MoLFormer	一种利用线性注意力和旋转位置编码技术，在SMILES 序列上通过掩码语言建模预训练的大规模 Transformer 模型。	完成	已验证	未验证	未验证	未验证	MoLFormer
ProtBERT	基于BERT架构，当前实现仅支持掩码语言模型（MLM）预训练和分类模式。	完成	已验证	未验证	未验证	未验证	ProtBERT
InfoGraph	一个用于无监督图级别表示学习的模型，通过互信息最大化（InfoMax）来学习分子图的嵌入表示。	完成	已验证	未验证	未验证	未验证	InfoGraph

后续我们会持续增加适配模型以及对应的样例notebook/代码。