#!/usr/bin/env python3
"""
Benchmark GNN models on standard datasets.
This script provides a simple way to benchmark different GNN architectures
on common datasets and compare their performance.
Usage:
python benchmark_model.py --models gcn gat --dataset Cora
python benchmark_model.py --models gcn --dataset Cora --epochs 200 --runs 10
"""
import argparse
import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, GATConv, SAGEConv, GINConv
from torch_geometric.datasets import Planetoid, TUDataset
from torch_geometric.loader import DataLoader
from torch_geometric.nn import global_mean_pool
import time
import numpy as np
class GCN(torch.nn.Module):
def __init__(self, num_features, hidden_channels, num_classes, dropout=0.5):
super().__init__()
self.conv1 = GCNConv(num_features, hidden_channels)
self.conv2 = GCNConv(hidden_channels, num_classes)
self.dropout = dropout
def forward(self, x, edge_index, batch=None):
x = self.conv1(x, edge_index)
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.conv2(x, edge_index)
if batch is not None:
x = global_mean_pool(x, batch)
return F.log_softmax(x, dim=1)
class GAT(torch.nn.Module):
def __init__(self, num_features, hidden_channels, num_classes, heads=8, dropout=0.6):
super().__init__()
self.conv1 = GATConv(num_features, hidden_channels, heads=heads, dropout=dropout)
self.conv2 = GATConv(hidden_channels * heads, num_classes, heads=1,
concat=False, dropout=dropout)
self.dropout = dropout
def forward(self, x, edge_index, batch=None):
x = F.dropout(x, p=self.dropout, training=self.training)
x = F.elu(self.conv1(x, edge_index))
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.conv2(x, edge_index)
if batch is not None:
x = global_mean_pool(x, batch)
return F.log_softmax(x, dim=1)
class GraphSAGE(torch.nn.Module):
def __init__(self, num_features, hidden_channels, num_classes, dropout=0.5):
super().__init__()
self.conv1 = SAGEConv(num_features, hidden_channels)
self.conv2 = SAGEConv(hidden_channels, num_classes)
self.dropout = dropout
def forward(self, x, edge_index, batch=None):
x = self.conv1(x, edge_index)
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.conv2(x, edge_index)
if batch is not None:
x = global_mean_pool(x, batch)
return F.log_softmax(x, dim=1)
MODELS = {
'gcn': GCN,
'gat': GAT,
'graphsage': GraphSAGE,
}
def train_node_classification(model, data, optimizer):
"""Train for node classification."""
model.train()
optimizer.zero_grad()
out = model(data.x, data.edge_index)
loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
loss.backward()
optimizer.step()
return loss.item()
@torch.no_grad()
def test_node_classification(model, data):
"""Test for node classification."""
model.eval()
out = model(data.x, data.edge_index)
pred = out.argmax(dim=1)
accs = []
for mask in [data.train_mask, data.val_mask, data.test_mask]:
correct = (pred[mask] == data.y[mask]).sum()
accs.append(float(correct) / int(mask.sum()))
return accs
def train_graph_classification(model, loader, optimizer, device):
"""Train for graph classification."""
model.train()
total_loss = 0
for data in loader:
data = data.to(device)
optimizer.zero_grad()
out = model(data.x, data.edge_index, data.batch)
loss = F.nll_loss(out, data.y)
loss.backward()
optimizer.step()
total_loss += loss.item() * data.num_graphs
return total_loss / len(loader.dataset)
@torch.no_grad()
def test_graph_classification(model, loader, device):
"""Test for graph classification."""
model.eval()
correct = 0
for data in loader:
data = data.to(device)
out = model(data.x, data.edge_index, data.batch)
pred = out.argmax(dim=1)
correct += (pred == data.y).sum().item()
return correct / len(loader.dataset)
def benchmark_node_classification(model_name, dataset_name, epochs, lr, weight_decay, device):
"""Benchmark a model on node classification."""
# Load dataset
dataset = Planetoid(root=f'/tmp/{dataset_name}', name=dataset_name)
data = dataset[0].to(device)
# Create model
model_class = MODELS[model_name]
model = model_class(
num_features=dataset.num_features,
hidden_channels=64,
num_classes=dataset.num_classes
).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
# Training
start_time = time.time()
best_val_acc = 0
best_test_acc = 0
for epoch in range(1, epochs + 1):
loss = train_node_classification(model, data, optimizer)
train_acc, val_acc, test_acc = test_node_classification(model, data)
if val_acc > best_val_acc:
best_val_acc = val_acc
best_test_acc = test_acc
train_time = time.time() - start_time
return {
'train_acc': train_acc,
'val_acc': best_val_acc,
'test_acc': best_test_acc,
'train_time': train_time,
}
def benchmark_graph_classification(model_name, dataset_name, epochs, lr, device):
"""Benchmark a model on graph classification."""
# Load dataset
dataset = TUDataset(root=f'/tmp/{dataset_name}', name=dataset_name)
# Split dataset
dataset = dataset.shuffle()
train_dataset = dataset[:int(len(dataset) * 0.8)]
test_dataset = dataset[int(len(dataset) * 0.8):]
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32)
# Create model
model_class = MODELS[model_name]
model = model_class(
num_features=dataset.num_features,
hidden_channels=64,
num_classes=dataset.num_classes
).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
# Training
start_time = time.time()
for epoch in range(1, epochs + 1):
loss = train_graph_classification(model, train_loader, optimizer, device)
# Final evaluation
train_acc = test_graph_classification(model, train_loader, device)
test_acc = test_graph_classification(model, test_loader, device)
train_time = time.time() - start_time
return {
'train_acc': train_acc,
'test_acc': test_acc,
'train_time': train_time,
}
def run_benchmark(args):
"""Run benchmark experiments."""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
# Determine task type
if args.dataset in ['Cora', 'CiteSeer', 'PubMed']:
task = 'node_classification'
else:
task = 'graph_classification'
print(f"\\nDataset: {args.dataset}")
print(f"Task: {task}")
print(f"Models: {', '.join(args.models)}")
print(f"Epochs: {args.epochs}")
print(f"Runs: {args.runs}")
print("=" * 60)
results = {model: [] for model in args.models}
# Run experiments
for run in range(args.runs):
print(f"\\nRun {run + 1}/{args.runs}")
print("-" * 60)
for model_name in args.models:
if model_name not in MODELS:
print(f"Unknown model: {model_name}")
continue
print(f" Training {model_name.upper()}...", end=" ")
try:
if task == 'node_classification':
result = benchmark_node_classification(
model_name, args.dataset, args.epochs,
args.lr, args.weight_decay, device
)
print(f"Test Acc: {result['test_acc']:.4f}, "
f"Time: {result['train_time']:.2f}s")
else:
result = benchmark_graph_classification(
model_name, args.dataset, args.epochs, args.lr, device
)
print(f"Test Acc: {result['test_acc']:.4f}, "
f"Time: {result['train_time']:.2f}s")
results[model_name].append(result)
except Exception as e:
print(f"Error: {e}")
# Print summary
print("\\n" + "=" * 60)
print("BENCHMARK RESULTS")
print("=" * 60)
for model_name in args.models:
if not results[model_name]:
continue
test_accs = [r['test_acc'] for r in results[model_name]]
times = [r['train_time'] for r in results[model_name]]
print(f"\\n{model_name.upper()}")
print(f" Test Accuracy: {np.mean(test_accs):.4f} ± {np.std(test_accs):.4f}")
print(f" Training Time: {np.mean(times):.2f} ± {np.std(times):.2f}s")
def main():
parser = argparse.ArgumentParser(description="Benchmark GNN models")
parser.add_argument('--models', nargs='+', default=['gcn'],
help='Model types to benchmark (gcn, gat, graphsage)')
parser.add_argument('--dataset', type=str, default='Cora',
help='Dataset name (Cora, CiteSeer, PubMed, ENZYMES, PROTEINS)')
parser.add_argument('--epochs', type=int, default=200,
help='Number of training epochs')
parser.add_argument('--runs', type=int, default=5,
help='Number of runs to average over')
parser.add_argument('--lr', type=float, default=0.01,
help='Learning rate')
parser.add_argument('--weight-decay', type=float, default=5e-4,
help='Weight decay for node classification')
args = parser.parse_args()
run_benchmark(args)
if __name__ == '__main__':
main()