liugangcode/torch-molecule

torch‑molecule – 分子发现的深度学习（sklearn风格API）

是什么 – 一个基于 PyTorch 的 Python 库，让研究人员和工程师能够使用熟悉的 scikit‑learn 风格接口来训练、评估和部署分子 AI 模型（预测、生成和表示模型）。

为何重要 – 分子发现（药物设计、材料科学、毒性预测等）越来越多地依赖图神经网络和 Transformer 模型。torch‑molecule 集成了大量最先进的架构、预训练检查点和基准数据集，同时隐藏了数据处理、模型构建和超参数调优的样板代码。

---

主要组件

组件	功能	示例模型
预测模型	对分子图或 SMILES 字符串进行监督回归/分类	GRIN, BFGNN, GREA, DIR, SGIR, GNNs, Transformer (SMILES), LSTM (SMILES)
生成模型	从头分子生成、条件生成、骨架引导采样	DeFoG, Graph DiT, DiGress, GDSS, MolGPT, JTVAE, GraphGA, 预训练生成器（NovoMolGen, MolGen‑large, Molexar, SAFE‑GPT）
表示模型	图或 SMILES 编码器的自监督预训练，用于下游任务	MoAMa, GraphMAE, AttrMasking, ContextPred, EdgePred, InfoGraph, 一系列预训练的 BERT/GPT 风格模型

内置数据集

• qm9 – 量子化学性质（约134k分子）
• chembl2k / broad6k – 药物样生物活性数据
• toxcast – 毒性标签
• admet – 吸收/分布/代谢/排泄/毒性
• gasperm – 聚合物气体渗透性
• zinc250k – 无标签的 ZINC 子集，用于无条件生成

所有数据集均可通过 torch_molecule.datasets.load_<name>() 加载，返回一个包含 .data（SMILES 列表）和 .target（numpy 数组或 None）的简单对象。

安装

# 创建新的 conda 环境（推荐 Python 3.11）
conda create -n torch_molecule python=3.11.7 && conda activate torch_molecule

# 从 PyPI 安装包
pip install torch-molecule

获取最新源码：

git clone https://github.com/liugangcode/torch-molecule
cd torch-molecule
pip install .

部分模型需要额外依赖（如 torch-scatter, transformers, selfies, fragment-selfies, molexar, safe-mol）。README 中列出了确切的 pip 命令。

典型工作流（预测示例）

from torch_molecule.datasets import load_qm9
from torch_molecule import GREAMolecularPredictor

# 1️⃣ 加载并分割数据
qm9 = load_qm9(local_dir='torchmol_data')
train, val = qm9.train_test_split(test_size=0.2, method='scaffold', seed=42)

# 2️⃣ 实例化模型（sklearn 风格）
model = GREAMolecularPredictor(
    num_task=1,
    task_type='regression',
    evaluate_higher_better=False,
    verbose='progress_bar'
)

# 3️⃣ 自动调优超参数并拟合
model.autofit(
    X_train=train.data, y_train=train.target,
    X_val=val.data,   y_val=val.target,
    n_trials=10
)

# 4️⃣ 对新 SMILES 进行预测
preds = model.predict(['C1=CC=CC=C1', 'CCO'])

相同模式适用于生成模型（HFPretrainedMolecularGenerator）和表示编码器。

检查点处理

• 保存到 Hugging Face – model.save_to_hf(repo_id='user/repo', task_id='qm9_grea', ...)
• 从 Hugging Face 加载 – model.load_from_hf(repo_id='user/repo', local_cache='.../GREA_qm9.pt')
• 本地保存/加载 – model.save_to_local('qm9_grea.pt') 和 model.load_from_local('qm9_grea.pt')

可扩展性

• 可通过贡献加载函数添加新数据集。
• 社区贡献者已集成多个预训练生成器（MolGen, Molexar, SAFE‑GPT）和数据集分割工具。
• 该库为开源，欢迎提交 Pull Requests。

---

简而言之 – torch‑molecule 为您提供了一个一站式、sklearn 风格的工具箱，用于现代图和 Transformer 基础的分子 AI，包含基准数据、预训练检查点和便捷的模型持久化功能。