mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-09-18 15:03:20 +08:00
7eff2b3d02
三个预训练分子编码器作为额外的 chemical token 接入（三选一，互斥）：
- MolFormer-XL 768d（SMILES 序列，11亿分子预训练）
- GraphMVP 300d（2D 图自监督，ICLR'22）
- GROVER 3200d（图 Transformer，NeurIPS'20）
实现：embedding 离线编码存 npy，模型内按 {smiles: vector} 查表注入，
不改动前向逻辑，默认全关、不影响既有实验。
- models.py: 新增三个 CHEM_KEYS_WITH_*、统一的 _load_offline_emb 查表、
proj_input_dims 按开关裁剪、forward 注入；子类同步转发参数
- nested_cv_optuna.py / pretrain.py: 新增 --{molformer,graphmvp,grover}-emb/-csv
- scripts_run/encode_embeddings/: 三个离线编码脚本 + README（含权重来源与
新旧环境兼容补丁说明）
- results/pretrained_encoders/: 五组 15 折结果 summary
(trials15/inner3/repeats3/seed42, 均带 external 预训练)
注：npy embedding 与预训练权重未入库（可由编码脚本复现）
44 lines
1.9 KiB
Markdown
44 lines
1.9 KiB
Markdown
# 离线预训练分子 embedding
|
||
|
||
三个预训练编码器作为额外的 chemical token 接入模型。embedding 离线生成、
|
||
存为 `.npy`（按 csv 行顺序对齐），训练时通过 `{smiles: vector}` 查表注入。
|
||
|
||
npy 文件与预训练权重不入库，运行下列脚本可复现。
|
||
|
||
## MolFormer (768d)
|
||
`ibm/MoLFormer-XL-both-10pct`，HuggingFace 直接加载。
|
||
注意：需 transformers>=4.57（依赖 `transformers.masking_utils`），
|
||
与项目 pin 的 4.45 冲突，建议装到独立目录后用 sys.path 隔离。
|
||
见 `encode_molformer.py`。
|
||
|
||
## GraphMVP (300d)
|
||
ICLR'22，5 层 GIN + OGB AtomEncoder/BondEncoder。
|
||
权重：https://github.com/chao1224/GraphMVP → Drive 文件夹
|
||
`GraphMVP_complate_features_for_regression/GraphMVP/pretraining_model.pth`
|
||
注意：新版 OGB 的 chirality 类别数为 5（checkpoint 为 4），加载时按形状逐层对齐。
|
||
见 `encode_graphmvp.py`。
|
||
|
||
## GROVER (3200d = atom 1600 + bond 1600)
|
||
NeurIPS'20，dualtrans (hidden 800, depth 6)。
|
||
代码：https://github.com/tencent-ailab/grover，权重 grover_base。
|
||
原项目基于 Python 3.6.8 / PyTorch 1.1，在现代环境需两处补丁：
|
||
1. `grover/util/utils.py` 的 `torch.load` 加 `weights_only=False`
|
||
2. `build_model` 前用 checkpoint 自带 args 补齐 current_args 缺失字段，
|
||
并为 dropout 等推理期字段提供默认值
|
||
见 `encode_grover.py`。
|
||
|
||
## 使用
|
||
```bash
|
||
python -m lnp_ml.modeling.pretrain main --device cuda \
|
||
--output-dir models/pretrain/<name> \
|
||
--<name>-emb data/interim/<name>_emb_external.npy \
|
||
--<name>-csv data/external/all_data_LiON.csv
|
||
|
||
python -m lnp_ml.modeling.nested_cv_optuna \
|
||
--device cuda --n-trials 15 --epochs-per-trial 30 --n-repeats 3 --seed 42 \
|
||
--<name>-emb data/interim/<name>_emb.npy \
|
||
--<name>-csv data/interim/internal.csv \
|
||
--init-from-pretrain models/pretrain/<name>/pretrain_delivery.pt \
|
||
--output-dir models/full/<name>_s42
|
||
```
|