AIDD快速开始EPT下游任务LBA/MSP/QM9微调与评估的保姆级实战本文带你完整跑通 AIDD 仓库中 EPTEquivariant Pretrained Transformer等变预训练 Transformer在 LBA、MSP、QM9 三大下游任务上的微调与评估从环境搭建、数据预处理到一行命令微调并教你看懂 RMSE、AUROC、MAE 指标。全程零踩坑。 什么是 EPT一个预训练模型搞定三类分子任务EPT是 AIDD 项目中的分子表征学习模型全称 An Equivariant Pretrained Transformer for Unified 3D Molecular Representation Learning。它的核心思路是先在多域分子数据小分子 GEOM/PCQM4M-v2 蛋白质 PDB/PDBBind上做一次去噪预训练得到一个通用的 3D 分子编码器然后在下游任务上只需轻量微调即可直接使用。本次实战聚焦三个经典下游任务任务类型预测目标评估指标训练配置LBA回归配体-受体结合亲和力neglog_affRMSE / PCC / SPCCEPT/configs/LBA/ID30/train.yamlMSP二分类两个配体哪个结合更好AUROCEPT/configs/MSP/train.yamlQM9回归分子量子化学性质homo、gap 等 12 项MAEEPT/configs/QM9/train.yaml三者共用的编码器都是XTransEncoderAct6 层、8 头、hidden_size 512 的等变 Transformer区别仅在任务头LBA/QM9 用回归头PredictorNNNModel/PredictorNNDModelMSP 用成对分类头PairClassifier见 EPT/configs/MSP/train.yaml。 3分钟快速上手安装 EPT 运行环境EPT 子项目位于仓库的EPT/目录环境定义在 EPT/env.yamlPython 3.9 PyTorch 2.1.2/CUDA 12.1 PyG rdkit fair-esm。git clone https://gitcode.com/BAAI/AIDD cd AIDD/EPT conda env create -f env.yaml conda activate EPT硬件说明下游微调单卡即可各任务 batch_size 为 16QM9 为 64只有多卡预训练才需要 8 卡。 获取 EPT 预训练权重微调的起点是一个多域预训练 checkpoint有两条路方式一推荐直接用官方预训练权重。官方 checkpoint 的下载地址见 EPT/README.md 中的 Pretrain on Multi-Domain Dataset 一节下载后记住它的本地路径下文记作pretrained_ckpt。方式二自己跑预训练需要 8 卡 预处理的 GEOM/PCQM4M-v2/PDB/PDBBind 数据耗时较长GPU0,1,2,3,4,5,6,7 bash execute/pretrain.sh configs/PreTrain/pretrain.yaml预训练脚本 EPT/execute/pretrain.sh 会在多卡时自动加torchrun分布式前缀训练配置为 EPT/configs/PreTrain/pretrain.yaml。️ 数据预处理把 LBA、MSP、QM9 数据备好EPT 使用 LMDB/mmap 格式存储分子数据。所有命令均在AIDD/EPT目录下执行预处理后的数据统一输出到./processed/目录也可从官方 figshare 获取全部预处理数据集链接见 EPT/README.md。① QM9最简单数据自动下载python -m scripts.process_data.process_QM9 --out_dir ./processed/QM9 --using_hydrogen --download脚本 EPT/scripts/process_data/process_QM9.py 会自动下载 GDB9 数据并切分为 train/valid/test 三份 mmap处理函数支持mu、alpha、homo、lumo、gap、r2、zpve、U0、U、H、G、Cv共 12 项性质。② LBA结合亲和力先获取官方原始数据LMDB 格式的 train/val/test 三个 split数据源见 README再执行python -m scripts.process_data.process_LBA --base_path splits_dir --out_dir ./processed/LBA/ID30 # 处理 ID60 序列相似度划分时把 out_dir 换成 ./processed/LBA/ID60③ MSP成对比较python -m scripts.process_data.process_MSP --base_path splits_dir --out_dir ./processed/MSP⚠️ 注意README 中 LBA 命令带了.py后缀process_LBA.pypython -m模块路径不要带.py上面命令才是正确写法。 核心实战一行命令微调 EPT 三大任务EPT 的微调入口是 EPT/train.py官方把训练 评估封装进了三个 shell 脚本只需传两个参数配置目录和预训练权重路径。脚本内部会通过pretrain_ckpt${CKPT}参数覆盖配置中的默认值见 EPT/train.py 的权重加载逻辑。LBA 结合亲和力微调GPU0 SPLITID30 bash execute/finetune_lba.sh ./configs/LBA pretrained_ckpt # 想用 60% 序列相似度划分把 SPLITID30 改成 SPLITID60SPLIT环境变量决定加载 EPT/configs/LBA/ID30/train.yaml 还是 EPT/configs/LBA/ID60/train.yaml默认 ID30见 EPT/execute/finetune_lba.sh。ID30/ID60 代表受体蛋白按 30%/60% 序列相似度划分ID60 更难、更考验泛化。MSP 二分类微调GPU0 bash execute/finetune_msp.sh ./configs/MSP pretrained_ckptMSP 任务把同一受体 两个配体打包成一条样本模型预测哪个结合更好正例 结合更强的配体配置见 EPT/configs/MSP/train.yaml最多训练 10 个 epoch、带早停patience 10。QM9 性质回归微调GPU0 PROPhomo bash execute/finetune_qm9.sh ./configs/QM9 pretrained_ckpt # 换性质PROPgap、PROPlumo 等 12 项性质任选PROP环境变量会同时注入dataset.train.property与dataset.valid.property见 EPT/execute/finetune_qm9.sh。能量类性质homo、lumo、gap、U、H、G 等会在数据集内自动换算成 eV见 EPT/data/dataset_qm9.py。训练产物在哪三个任务的最优 checkpoint 分别保存在./ckpts/LBA/ID30|ID60、./ckpts/MSP、./ckpts/QM9的checkpoint/子目录中topk_map.txt记录各 top-k 权重的验证指标。 评估指标解读RMSE、AUROC、MAE 告诉你什么三个 finetune 脚本在训练结束后会自动执行评估如 EPT/execute/finetune_lba.sh 最后一行手动重跑命令如下# LBA输出 RMSE、PCC、SPCC python -m evaluation.eval_lba --config ./configs/LBA/ID30/test.yaml --ckpt ./ckpts/LBA/ID30 --gpu 0 # MSP输出 AUROC python -m evaluation.eval_msp --config ./configs/MSP/test.yaml --ckpt ./ckpts/MSP --gpu 0 # QM9输出 MAE注意补充 property 覆盖参数 python -m evaluation.eval_qm9 --config ./configs/QM9/test.yaml --ckpt ./ckpts/QM9 --gpu 0 \ dataset.train.propertyhomo dataset.test.propertyhomo指标怎么读指标出现任务含义优劣RMSELBA结合亲和力预测均方根误差kcal/mol 尺度越小越好PCC / SPCCLBA皮尔逊 / 斯皮尔曼相关系数越接近 1 越好AUROCMSP成对排序的受试者工作特征曲线下面积越接近 1 越好MAEQM9性质预测平均绝对误差eV 尺度越小越好细节亮点LBA 评估会对 top-k 个 checkpoint 做权重平均后再推理EPT/evaluation/eval_lba.py通常能进一步压低 RMSEQM9 评估使用 MAD 归一化后再还原量纲EPT/evaluation/eval_qm9.py。⚠️ 新手常见坑清单QM9 的 test.yaml 是坑王EPT/configs/QM9/test.yaml 中mmap_dir硬编码了作者机器的绝对路径/GLOBALFS/...运行前必须改成你本地路径dataset: train: mmap_dir: ./processed/QM9/atomH/train test: mmap_dir: ./processed/QM9/atomH/test否则eval_qm9会直接报FileNotFoundError。QM9 默认训练 1000 个 epochEPT/configs/QM9/train.yaml 中max_epoch: 1000且patience: -1不早停。资源紧张时可手动把max_epoch调小或直接用命令行覆盖python train.py --gpus 0 --config configs/QM9/train.yaml pretrain_ckptckpt \ dataset.train.propertyhomo dataset.valid.propertyhomo trainer.config.max_epoch50pretrain_ckpt必须是脚本第二个参数finetune 脚本只认$1配置目录和$2权重路径传反或不传都会加载null权重、退化成从零训练。环境变量要写在命令同一行GPU0 SPLITID60 bash execute/finetune_lba.sh ...这种写法才生效分两行 export 也可但注意SPLIT/PROP只对 LBA/QM9 有意义。GPU默认是 -1纯 CPU不写GPU0时脚本不会报错只是慢得离谱新手容易以为是程序卡死。LBA/MSP 原始数据需要联网下载LMDB 原始 split 来自官方数据源各数据集下载链接见 EPT/README.md 的 Downstream Data 一节QM9 则是--download参数自动拉取。 EPT 相关文件索引类别文件微调入口EPT/train.py执行脚本EPT/execute/finetune_lba.sh、EPT/execute/finetune_msp.sh、EPT/execute/finetune_qm9.sh、EPT/execute/pretrain.sh配置文件EPT/configs/LBA/、EPT/configs/MSP/、EPT/configs/QM9/、EPT/configs/PreTrain/pretrain.yaml数据预处理EPT/scripts/process_data/process_LBA.py、EPT/scripts/process_data/process_MSP.py、EPT/scripts/process_data/process_QM9.py数据集实现EPT/data/dataset_lba.py、EPT/data/dataset_msp.py、EPT/data/dataset_qm9.py评估脚本EPT/evaluation/eval_lba.py、EPT/evaluation/eval_msp.py、EPT/evaluation/eval_qm9.py模型封装EPT/models/wrappers/predictor / classifier / denoise 等任务头 小结到这里你已经走完了 EPT 的完整下游工作流装环境 → 拿权重 → 备数据 → 一行命令微调 → 自动出指标。相比从零搭建分子模型EPT 的价值在于预训练一次、三任务复用——同一套等变 Transformer 编码器换不同任务头就能在结合亲和力回归LBA、成对结合比较MSP、量子化学性质回归QM9上快速落地。如果想深入模型结构可以从 EPT/models/TransAllAtom/ 的编码器实现和 EPT/models/graph_constructor.py 的图构建逻辑读起。若该工作对你的研究有帮助引用请参考 EPT/README.md 中的 Nature Communications (2026) 文献信息。创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考