从ESOL水溶性预测出发：如何为你的分子属性预测任务定制PyTorch Geometric数据集-平芜编程栈

从ESOL水溶性预测到分子图数据集构建：PyTorch Geometric实战指南

在药物发现和材料设计领域，分子属性预测一直是核心挑战之一。传统方法依赖耗时费力的实验测量，而图神经网络(GNN)的出现为这一领域带来了新的可能性。PyTorch Geometric作为领先的图深度学习框架，其内置的MoleculeNet数据集虽然方便，但面对真实研究场景时，我们往往需要从零构建自己的分子图数据集。本文将深入解析如何将SMILES字符串转化为PyTorch Geometric可处理的图数据结构，并提供一个可复用的完整流程模板。

1. 分子图数据基础：从SMILES到图表示

SMILES(Simplified Molecular Input Line Entry System)字符串是描述分子结构的标准化学标识符。例如，阿司匹林的SMILES表示为CC(=O)OC1=CC=CC=C1C(=O)O。这种线性表示需要转换为图结构才能被GNN处理，其中原子作为节点，化学键作为边。

关键转换步骤：

原子特征提取：每个原子节点需要编码多种属性：

atom_features = [ atomic_number, # 原子序数 chiral_tag, # 手性信息 total_degree, # 连接度 formal_charge, # 形式电荷 num_hs, # 连接氢原子数 hybridization, # 杂化类型 is_aromatic, # 芳香性 is_in_ring # 是否在环上 ]

键特征编码：每条边需要包含键类型等信息：

bond_features = [ bond_type, # 单键、双键等 bond_stereo, # 立体化学 is_conjugated # 是否共轭 ]

提示：RDKit库提供了完整的化学信息处理功能，是SMILES解析的核心工具

2. 构建自定义分子数据集的完整流程

2.1 原始数据准备与解析

典型的分子数据集通常以CSV格式存储，包含SMILES字符串和对应的属性标签。我们需要：

读取并解析原始文件：

import pandas as pd raw_data = pd.read_csv('custom_molecules.csv') smiles_list = raw_data['smiles'].tolist() labels = raw_data['log_solubility'].values

数据清洗与验证：
- 检查SMILES有效性
- 处理缺失值
- 标准化标签范围

2.2 自定义Dataset类实现

PyTorch Geometric通过InMemoryDataset类提供高效的数据加载机制。以下是核心实现步骤：

from torch_geometric.data import InMemoryDataset, Data import torch class MolecularDataset(InMemoryDataset): def __init__(self, root, smiles_list, labels, transform=None): self.smiles_list = smiles_list self.labels = labels super().__init__(root, transform) self.data, self.slices = torch.load(self.processed_paths[0]) @property def raw_file_names(self): return ['custom_data.csv'] @property def processed_file_names(self): return ['processed_data.pt'] def process(self): data_list = [] for smiles, label in zip(self.smiles_list, self.labels): mol = Chem.MolFromSmiles(smiles) if mol is None: # 跳过无效SMILES continue # 原子特征矩阵 x = self._get_node_features(mol) # 边索引和边特征 edge_index, edge_attr = self._get_edge_features(mol) data = Data( x=x, edge_index=edge_index, edge_attr=edge_attr, y=torch.tensor([label], dtype=torch.float), smiles=smiles ) data_list.append(data) data, slices = self.collate(data_list) torch.save((data, slices), self.processed_paths[0])

2.3 特征工程关键参数解析

from_smiles函数中的参数对图结构有重大影响：

参数	类型	默认值	影响
`with_hydrogen`	bool	False	是否显式包含氢原子
`kekulize`	bool	False	是否将芳香键转换为凯库勒式
`use_chirality`	bool	False	是否包含手性信息
`use_bond_order`	bool	True	是否考虑键级信息

选择建议：

对小分子体系，建议启用with_hydrogen
预测芳香性相关属性时，保持kekulize=False
手性敏感的分子系统需要设置use_chirality=True

3. 高级数据处理技巧

3.1 数据增强策略

分子图数据增强能有效提升模型泛化能力：

旋转不变性增强：随机旋转分子坐标（如有3D信息）
键扰动：在合理范围内调整键长/键角
原子掩码：随机遮蔽部分原子特征

def augment_molecule(data, p=0.1): # 原子特征掩码 if random.random() < p: mask = torch.rand(data.x.size(0)) < 0.1 data.x[mask] = 0 # 边特征扰动 if random.random() < p and data.edge_attr is not None: noise = torch.randn_like(data.edge_attr) * 0.05 data.edge_attr += noise return data

3.2 非标准分子处理

特殊分子结构需要额外处理：

金属有机框架(MOFs)：
- 考虑配位键的特殊表示
- 添加周期性边界条件处理
大分子系统：
- 实现分块处理策略
- 使用层次化图表示
互变异构体：
- 标准化互变异构形式
- 考虑多重表示集成

4. 数据集分割与评估策略

4.1 科学的数据分割方法

分子数据集分割需要考虑化学空间的覆盖：

方法	描述	适用场景
随机分割	简单随机划分	初步验证
骨架分割	按分子骨架分类	评估泛化性
时间分割	按发现时间划分	实际应用模拟
属性分割	按目标属性分层	不平衡数据

骨架分割实现示例：

from rdkit.Chem import AllChem def split_by_scaffold(dataset, test_ratio=0.2): scaffolds = {} for idx, data in enumerate(dataset): mol = Chem.MolFromSmiles(data.smiles) scaffold = AllChem.MurckoDecompose(mol) scaffold_smiles = Chem.MolToSmiles(scaffold) scaffolds.setdefault(scaffold_smiles, []).append(idx) scaffold_groups = list(scaffolds.values()) random.shuffle(scaffold_groups) split_idx = int((1 - test_ratio) * len(scaffold_groups)) train_indices = sum(scaffold_groups[:split_idx], []) test_indices = sum(scaffold_groups[split_idx:], []) return train_indices, test_indices

4.2 评估指标选择

不同分子属性预测任务需要针对性的评估指标：

任务类型	推荐指标	注意事项
回归任务	RMSE, R²	注意量纲一致性
分类任务	ROC-AUC, PR-AUC	处理类别不平衡
多任务	加权平均得分	任务权重分配

注意：分子属性预测中，不同指标可能给出矛盾结论，建议结合多个指标综合评估

5. 实战案例：从ESOL到自定义数据集

让我们通过一个完整案例演示如何构建水溶性预测数据集：

数据准备：

# 假设我们有自定义数据文件 custom_data = [ ("CC(=O)OC1=CC=CC=C1C(=O)O", -0.72), # 阿司匹林 ("C1=CC=C(C=C1)C=O", -0.53), # 苯甲醛 ("CCO", 0.21) # 乙醇 ] df = pd.DataFrame(custom_data, columns=["smiles", "logS"]) df.to_csv("custom_solubility.csv", index=False)

数据集实现：

dataset = MolecularDataset( root="custom_data", smiles_list=df["smiles"].tolist(), labels=df["logS"].values ) # 数据集使用示例 train_loader = DataLoader(dataset[:800], batch_size=32, shuffle=True) val_loader = DataLoader(dataset[800:], batch_size=32)

模型训练验证：

model = GNN(in_features=dataset.num_features) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) for epoch in range(100): model.train() for batch in train_loader: optimizer.zero_grad() pred = model(batch.x, batch.edge_index, batch.batch) loss = F.mse_loss(pred, batch.y) loss.backward() optimizer.step() # 验证步骤 model.eval() val_loss = 0 for batch in val_loader: with torch.no_grad(): pred = model(batch.x, batch.edge_index, batch.batch) val_loss += F.mse_loss(pred, batch.y) print(f"Epoch {epoch}, Val Loss: {val_loss/len(val_loader):.4f}")

在实际项目中，我们还需要考虑特征标准化、模型正则化、早停等技巧。一个完整的分子属性预测系统通常需要多次迭代优化数据表示和模型架构。