企业官网建设流程全解析

AIDD快速开始：EPT下游任务LBA/MSP/QM9微调与评估的保姆级实战

本文带你完整跑通 AIDD 仓库中 EPT（Equivariant 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 分子编码器，然后在下游任务上只需轻量微调即可直接使用。

本次实战聚焦三个经典下游任务：

三者共用的编码器都是XTransEncoderAct（6 层、8 头、hidden_size 512 的等变 Transformer），区别仅在任务头：LBA/QM9 用回归头PredictorNNNModel/PredictorNNDModel，MSP 用成对分类头PairClassifier（见 EPT/configs/MSP/train.yaml）。

🧰 3分钟快速上手：安装 EPT 运行环境

EPT 子项目位于仓库的EPT/目录，环境定义在 EPT/env.yaml（Python 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 为 16，QM9 为 64）；只有多卡预训练才需要 8 卡。

📦 获取 EPT 预训练权重

微调的起点是一个多域预训练 checkpoint，有两条路：

方式一（推荐）：直接用官方预训练权重。官方 checkpoint 的下载地址见 EPT/README.md 中的 "Pretrain on Multi-Domain Dataset" 一节，下载后记住它的本地路径（下文记作<pretrained_ckpt>）。

方式二：自己跑预训练（需要 8 卡 + 预处理的 GEOM/PCQM4M-v2/PDB/PDBBind 数据，耗时较长）：

GPU=0,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.py），python -m模块路径不要带.py，上面命令才是正确写法。

🚀 核心实战：一行命令微调 EPT 三大任务

EPT 的微调入口是 EPT/train.py，官方把"训练 + 评估"封装进了三个 shell 脚本，只需传两个参数：<配置目录>和<预训练权重路径>。脚本内部会通过pretrain_ckpt=${CKPT}参数覆盖配置中的默认值（见 EPT/train.py 的权重加载逻辑）。

LBA 结合亲和力微调

GPU=0 SPLIT=ID30 bash execute/finetune_lba.sh ./configs/LBA <pretrained_ckpt> # 想用 60% 序列相似度划分？把 SPLIT=ID30 改成 SPLIT=ID60

SPLIT环境变量决定加载 EPT/configs/LBA/ID30/train.yaml 还是 EPT/configs/LBA/ID60/train.yaml，默认 ID30（见 EPT/execute/finetune_lba.sh）。ID30/ID60 代表受体蛋白按 30%/60% 序列相似度划分，ID60 更难、更考验泛化。

MSP 二分类微调

GPU=0 bash execute/finetune_msp.sh ./configs/MSP <pretrained_ckpt>

MSP 任务把"同一受体 + 两个配体"打包成一条样本，模型预测哪个结合更好（正例 = 结合更强的配体），配置见 EPT/configs/MSP/train.yaml，最多训练 10 个 epoch、带早停（patience 10）。

QM9 性质回归微调

GPU=0 PROP=homo bash execute/finetune_qm9.sh ./configs/QM9 <pretrained_ckpt> # 换性质：PROP=gap、PROP=lumo 等 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.property=homo dataset.test.property=homo

指标怎么读：

细节亮点：LBA 评估会对 top-k 个 checkpoint 做权重平均后再推理（EPT/evaluation/eval_lba.py），通常能进一步压低 RMSE；QM9 评估使用 MAD 归一化后再还原量纲（EPT/evaluation/eval_qm9.py）。

⚠️ 新手常见坑清单

1. 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。

2. QM9 默认训练 1000 个 epoch：EPT/configs/QM9/train.yaml 中max_epoch: 1000且patience: -1（不早停）。资源紧张时可手动把max_epoch调小，或直接用命令行覆盖：

   python train.py --gpus 0 --config configs/QM9/train.yaml pretrain_ckpt=<ckpt> \ dataset.train.property=homo dataset.valid.property=homo trainer.config.max_epoch=50

3. pretrain_ckpt必须是脚本第二个参数：finetune 脚本只认$1（配置目录）和$2（权重路径），传反或不传都会加载null权重、退化成从零训练。

4. 环境变量要写在命令同一行：GPU=0 SPLIT=ID60 bash execute/finetune_lba.sh ...这种写法才生效；分两行 export 也可，但注意SPLIT/PROP只对 LBA/QM9 有意义。

5. GPU默认是 -1（纯 CPU）：不写GPU=0时脚本不会报错，只是慢得离谱，新手容易以为是程序卡死。

6. LBA/MSP 原始数据需要联网下载：LMDB 原始 split 来自官方数据源（各数据集下载链接见 EPT/README.md 的 "Downstream Data" 一节），QM9 则是--download参数自动拉取。

📂 EPT 相关文件索引

🔚 小结

到这里，你已经走完了 EPT 的完整下游工作流：装环境 → 拿权重 → 备数据 → 一行命令微调 → 自动出指标。相比从零搭建分子模型，EPT 的价值在于"预训练一次、三任务复用"——同一套等变 Transformer 编码器，换不同任务头就能在结合亲和力回归（LBA）、成对结合比较（MSP）、量子化学性质回归（QM9）上快速落地。如果想深入模型结构，可以从 EPT/models/TransAllAtom/ 的编码器实现和 EPT/models/graph_constructor.py 的图构建逻辑读起。

若该工作对你的研究有帮助，引用请参考 EPT/README.md 中的 Nature Communications (2026) 文献信息。