#!/usr/bin/env python3"""Benchmark GNN models on standard datasets. This script provides a simple way to benchmark different GNN architectureson 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 argparseimport torchimport torch.nn.functional as Ffrom torch_geometric.nn import GCNConv, GATConv, SAGEConv, GINConvfrom torch_geometric.datasets import Planetoid, TUDatasetfrom torch_geometric.loader import DataLoaderfrom torch_geometric.nn import global_mean_poolimport timeimport 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()