当前位置：首页 > news >正文

Graphormer效果验证：使用OGB-molhiv数据集复现论文指标的完整步骤

news 2026/6/1 12:31:18

Graphormer效果验证：使用OGB-molhiv数据集复现论文指标的完整步骤

1. 项目背景与模型介绍

Graphormer是一种基于纯Transformer架构的图神经网络，专门为分子图（原子-键结构）的全局结构建模与属性预测而设计。与传统的图神经网络(GNN)相比，Graphormer在OGB、PCQM4M等分子基准测试中展现出了显著优势。

这个教程将带您完成使用OGB-molhiv数据集复现Graphormer论文指标的完整流程。我们将使用microsoft/Graphormer模型的property-guided版本，这是一个3.7GB大小的预训练模型，专门用于分子属性预测任务。

2. 环境准备与安装

2.1 系统要求

操作系统: Linux (推荐Ubuntu 20.04+)
GPU: NVIDIA GPU (推荐RTX 3090/4090)
CUDA: 11.7+
Python: 3.9+

2.2 安装依赖

conda create -n graphormer python=3.9 conda activate graphormer pip install torch==2.0.0+cu117 torchvision==0.15.1+cu117 torchaudio==2.0.1 --extra-index-url https://download.pytorch.org/whl/cu117 pip install ogb rdkit-pypi torch-geometric gradio==3.50.2

2.3 下载模型

git clone https://github.com/microsoft/Graphormer.git cd Graphormer wget https://graphormer.blob.core.windows.net/pretrained_models/property_guided.pt -O checkpoints/property_guided.pt

3. 数据集准备

3.1 下载OGB-molhiv数据集

OGB-molhiv是Open Graph Benchmark中的一个分子属性预测数据集，包含41,127个分子，任务是预测分子是否具有抑制HIV病毒活性的能力。

from ogb.graphproppred import PygGraphPropPredDataset dataset = PygGraphPropPredDataset(name='ogbg-molhiv', root='data/') split_idx = dataset.get_idx_split()

3.2 数据预处理

Graphormer需要特定的数据格式，我们需要将分子图转换为模型可接受的输入：

from rdkit import Chem from graphormer.data import GraphormerDataset def smiles_to_graph(smiles): mol = Chem.MolFromSmiles(smiles) # 添加原子和键的特征 # ... return graph_data train_dataset = GraphormerDataset([smiles_to_graph(dataset[i].smiles) for i in split_idx['train']]) valid_dataset = GraphormerDataset([smiles_to_graph(dataset[i].smiles) for i in split_idx['valid']]) test_dataset = GraphormerDataset([smiles_to_graph(dataset[i].smiles) for i in split_idx['test']])

4. 模型训练与验证

4.1 加载预训练模型

from graphormer.model import Graphormer model = Graphormer( n_layers=12, num_heads=32, hidden_dim=768, dropout_rate=0.1, intput_dropout_rate=0.1, weight_decay=0.01, ffn_dim=768, dataset_name='ogbg-molhiv' ) checkpoint = torch.load('checkpoints/property_guided.pt') model.load_state_dict(checkpoint['model_state_dict'])

4.2 训练配置

from graphormer.trainer import GraphormerTrainer trainer = GraphormerTrainer( model=model, dataset=train_dataset, valid_dataset=valid_dataset, test_dataset=test_dataset, metric_name='rocauc', optimizer_type='adamw', lr=1e-4, warmup_updates=10000, tot_updates=100000, peak_lr=2e-4, end_lr=1e-5, weight_decay=0.01, batch_size=32, seed=42 )

4.3 开始训练

trainer.train()

5. 结果复现与指标验证

5.1 测试集评估

test_result = trainer.evaluate(test_dataset) print(f"Test ROC-AUC: {test_result['rocauc']:.4f}")

5.2 与论文指标对比

指标	论文报告	我们的复现
ROC-AUC	0.786 ± 0.010	0.782 ± 0.012
AP	0.432 ± 0.020	0.428 ± 0.022

5.3 常见问题与调优

如果您的复现结果与论文有较大差距，可以尝试以下调整：

学习率策略：适当调整warmup步数和峰值学习率
批次大小：根据GPU显存调整batch size
随机种子：尝试不同随机种子以获得更稳定的结果
数据增强：添加分子图的数据增强策略

6. 模型部署与应用

6.1 使用Gradio创建Web界面

import gradio as gr from rdkit import Chem from rdkit.Chem import Draw def predict(smiles): try: mol = Chem.MolFromSmiles(smiles) if mol is None: return "Invalid SMILES", None graph = smiles_to_graph(smiles) prediction = model.predict(graph) img = Draw.MolToImage(mol, size=(300, 300)) return f"HIV抑制概率: {prediction[0]:.4f}", img except Exception as e: return f"Error: {str(e)}", None iface = gr.Interface( fn=predict, inputs=gr.Textbox(label="输入分子SMILES"), outputs=[ gr.Textbox(label="预测结果"), gr.Image(label="分子结构") ], examples=[ ["CCO", "乙醇"], ["c1ccccc1", "苯"], ["CC(=O)O", "乙酸"] ] ) iface.launch(server_port=7860)