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

ChemBERTa模型迁移适配指导

1. 模型概述

ChemBERTa是一种基于Transformer架构的深度学习模型，专门用于处理 SMILES 字符串（简化分子线性输入系统），从而实现对分子结构的高效表示与性质预测。该模型以 RoBERTa 为基础架构，结合化学领域的预训练策略，在药物发现、材料科学和计算化学等任务中展现出强大能力，本文描述的ChemBERTa模型是基于DeepChem套件实现的，后续适配也是基于该套件修改。

2. 准备运行环境

2.1 软件环境

组件版本
Python3.10.19
PyTorch2.1.0
torch_npu2.1.0.post13
CANN8.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=chemberta
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 deepchem --clone PyTorch-2.1.0
conda activate deepchem

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 os
import tempfile
import shutil
import pandas as pd
import deepchem as dc

tempdir = tempfile.mkdtemp()

# preparing dataset
smiles = ["CCN(CCSC)C(=O)N[C@@](C)(CC)C(F)(F)F","CC1(C)CN(C(=O)Nc2cc3ccccc3nn2)C[C@@]2(CCOC2)O1"]
labels = [3.112,2.432]
df = pd.DataFrame(list(zip(smiles, labels)), columns=["smiles", "task1"])
with dc.utils.UniversalNamedTemporaryFile(mode='w') as tmpfile:
    df.to_csv(tmpfile.name)
    loader = dc.data.CSVLoader(["task1"], feature_field="smiles", featurizer=dc.feat.DummyFeaturizer())
    dataset = loader.create_dataset(tmpfile.name)
# Pretraining in MLM mode
from deepchem.models.torch_models.chemberta import Chemberta
pretrain_model_dir = os.path.join(tempdir, 'pretrain-model')

tokenizer_path = 'seyonec/PubChem10M_SMILES_BPE_60k'
pretrain_model = Chemberta(task='mlm', model_dir=pretrain_model_dir, tokenizer_path=tokenizer_path)
pretrain_loss = pretrain_model.fit(dataset, nb_epoch=1)

# finetuning in regression mode
finetune_model_dir = os.path.join(tempdir, 'finetune-model')
finetune_model = Chemberta(task='regression', model_dir=finetune_model_dir, tokenizer_path=tokenizer_path)
finetune_model.load_from_pretrained(pretrain_model_dir)
finetuning_loss = finetune_model.fit(dataset, nb_epoch=1)

print(f"模型所在设备: {finetune_model.device}")

# prediction and evaluation
result = finetune_model.predict(dataset)
eval_results = finetune_model.evaluate(dataset, metrics=dc.metrics.Metric(dc.metrics.mae_score))
print(eval_results)
# removing temporary directory
if os.path.exists(tempdir):
    shutil.rmtree(tempdir)

3.5 运行测试代码

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

export HF_ENDPOINT=https://hf-mirror.com
pytest test_chemberta.py

测试结果： image