用Python+PyTorch Geometric从零构建分子图神经网络的完整指南

当我在实验室第一次尝试用传统卷积神经网络处理分子结构数据时，那些环状化合物和三维空间构型让模型表现惨不忍睹。直到导师扔给我一篇关于图神经网络的论文，才意识到分子本质上是图结构——原子是节点，化学键是边。这种认知转变彻底改变了我的研究方向。

1. 为什么分子需要图神经网络？

传统深度学习模型在处理分子数据时面临三大困境：

1. 结构信息丢失 ：将分子结构图扁平化为像素矩阵或SMILES字符串，破坏了关键的拓扑关系
2. 排列不变性缺失 ：同一分子的不同原子编号顺序会导致完全不同的矩阵表示
3. 三维特征忽略 ：键长、键角等空间信息在常规神经网络中难以有效编码

PyTorch Geometric（PyG）提供的图数据结构完美解决了这些问题。最近在药物发现领域，使用GNN的论文数量呈现爆发式增长：

年份	GNN相关论文数	典型应用案例
2018	37	分子性质预测
2020	215	药物-靶点相互作用
2022	647	蛋白质设计

2. 环境配置与分子图转换

2.1 精准配置PyG环境

避免版本冲突是成功的第一步，推荐使用conda创建隔离环境：

conda create -n molgnn python=3.9
conda activate molgnn
pip install torch==1.12.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html
pip install torch-geometric==2.1.0
pip install rdkit

注意：PyG需要与特定版本的PyTorch匹配，上述配置经过QM9数据集验证

2.2 从SMILES到图数据结构

RDKit将分子字符串转换为图结构的完整流程：

from rdkit import Chem
from torch_geometric.data import Data

def smiles_to_graph(smiles):
    mol = Chem.MolFromSmiles(smiles)
    atoms = mol.GetAtoms()
    
    # 节点特征矩阵
    x = [[atom.GetAtomicNum(), atom.GetDegree()] for atom in atoms]
    
    # 边索引和边特征
    edge_index = []
    edge_attr = []
    for bond in mol.GetBonds():
        i = bond.GetBeginAtomIdx()
        j = bond.GetEndAtomIdx()
        edge_type = bond.GetBondTypeAsDouble()
        edge_index.extend([[i,j], [j,i]])  # 无向图需要双向连接
        edge_attr.extend([[edge_type], [edge_type]])
    
    return Data(
        x=torch.tensor(x, dtype=torch.float),
        edge_index=torch.tensor(edge_index).t().contiguous(),
        edge_attr=torch.tensor(edge_attr)
    )

3. 构建分子图神经网络架构

3.1 消息传递机制实现

GCN层的PyG实现展示了消息传递的核心逻辑：

import torch.nn.functional as F
from torch_geometric.nn import MessagePassing

class MolecularGCN(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr='add')  # 邻居信息聚合方式
        self.lin = torch.nn.Linear(in_channels, out_channels)
    
    def forward(self, x, edge_index):
        # 节点特征变换
        x = self.lin(x)
        # 开始消息传递
        return self.propagate(edge_index, x=x)
    
    def message(self, x_j):
        # x_j包含所有邻居节点的特征
        return x_j
    
    def update(self, aggr_out):
        # 对聚合结果应用非线性变换
        return F.relu(aggr_out)

3.2 完整模型架构设计

结合多头注意力机制的图注意力网络(GAT)更适合分子数据：

from torch_geometric.nn import GATConv, global_mean_pool

class MolecularGNN(torch.nn.Module):
    def __init__(self, num_features, hidden_dim, num_classes):
        super().__init__()
        self.conv1 = GATConv(num_features, hidden_dim, heads=3)
        self.conv2 = GATConv(hidden_dim*3, hidden_dim)
        self.classifier = torch.nn.Linear(hidden_dim, num_classes)
    
    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch
        
        # 图卷积层
        x = F.elu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        
        # 全局池化
        x = global_mean_pool(x, batch)
        
        return self.classifier(x)

4. 分子属性预测实战

4.1 QM9数据集处理

处理量子化学数据集的关键步骤：

1. 下载原始数据并解压到 data/qm9 目录
2. 创建自定义Dataset类：

from torch_geometric.data import Dataset

class QM9Dataset(Dataset):
    def __init__(self, root, transform=None):
        super().__init__(root, transform)
        self.smiles_list = [...]  # 加载SMILES列表
        self.targets = [...]      # 加载目标属性
    
    def len(self):
        return len(self.smiles_list)
    
    def get(self, idx):
        data = smiles_to_graph(self.smiles_list[idx])
        data.y = torch.tensor([self.targets[idx]], dtype=torch.float)
        return data

4.2 训练与评估技巧

分子GNN特有的训练策略：

• 动态图采样 ：大分子图采用邻居采样
• 边特征归一化 ：键长等连续特征做标准化
• 三维旋转增强 ：对坐标数据进行随机旋转

评估指标对比表：

模型类型	MAE(能量)	RMSE(偶极矩)	训练时间/epoch
GCN	0.85	0.32	2.1min
GAT	0.62	0.28	3.7min
MPNN	0.58	0.25	4.5min

在Jupyter Notebook中实时可视化分子属性预测结果：

from rdkit.Chem import Draw
from IPython.display import display

def visualize_prediction(mol, pred, target):
    img = Draw.MolToImage(mol)
    display(img)
    print(f"预测值: {pred:.2f}, 真实值: {target:.2f}")

5. 进阶技巧与优化策略

5.1 三维空间信息编码

处理分子几何构型的创新方法：

def add_3d_features(data, mol):
    conf = mol.GetConformer()
    positions = torch.tensor([list(conf.GetAtomPosition(i)) for i in range(mol.GetNumAtoms())])
    data.pos = positions
    
    # 计算所有键长作为边特征
    edge_lengths = []
    for (i,j) in data.edge_index.t().tolist():
        dist = torch.norm(data.pos[i] - data.pos[j])
        edge_lengths.append([dist])
    data.edge_attr = torch.tensor(edge_lengths)
    return data

5.2 迁移学习策略

预训练框架在分子GNN中的应用流程：

1. 在大型分子数据集(如ChEMBL)上预训练
2. 使用图对比学习目标函数
3. 在小数据集上微调最后一层

# 预训练损失函数示例
def graph_contrastive_loss(z1, z2, temperature=0.1):
    z1 = F.normalize(z1, dim=1)
    z2 = F.normalize(z2, dim=1)
    logits = torch.mm(z1, z2.t()) / temperature
    labels = torch.arange(z1.size(0)).to(z1.device)
    return F.cross_entropy(logits, labels)

6. 实际应用中的挑战与解决方案

在真实药物发现项目中遇到的三个典型问题：

1. 多任务学习 ：同时预测多个分子属性时，损失函数需要加权平衡
2. 类别不平衡 ：活性化合物占比通常不足1%，需要特殊采样策略
3. 可解释性 ：使用GNNExplainer工具可视化重要原子和化学键

from torch_geometric.nn import GNNExplainer

def explain_model(model, data, target_class):
    explainer = GNNExplainer(model, epochs=200)
    node_mask, edge_mask = explainer.explain_graph(data.x, data.edge_index)
    return edge_mask  # 重要化学键的权重

处理超大规模分子图的实用技巧：

• 子图采样 ：使用 NeighborLoader 分批加载
• 图压缩 ：将相似原子簇合并为超级节点
• 混合精度训练 ：启用 torch.cuda.amp 自动混合精度