Pytorch Geometric Graphsage Tutorial

PyTorch Geometric is a popular open-source library for geometric deep learning, which provides a wide range of tools and functionalities for working with graph-structured data. One of the key features of PyTorch Geometric is its support for GraphSAGE, a popular graph neural network (GNN) architecture. In this tutorial, we will provide a comprehensive overview of how to use PyTorch Geometric to implement GraphSAGE models for node classification tasks.

Introduction to GraphSAGE

GraphSAGE is a graph neural network architecture that uses a neighbor-sampling approach to learn node representations in a graph. The key idea behind GraphSAGE is to sample a fixed number of neighbors for each node and then use these neighbors to compute the node’s representation. This approach allows GraphSAGE to scale to large graphs and is particularly useful for node classification tasks.

PyTorch Geometric Implementation of GraphSAGE

PyTorch Geometric provides a pre-built implementation of GraphSAGE, which can be used to easily implement GraphSAGE models for node classification tasks. The PyTorch Geometric implementation of GraphSAGE is based on the GraphSAGE layer, which takes in a graph and outputs a node representation matrix. The GraphSAGE layer can be used in combination with other PyTorch Geometric layers, such as the GCNConv layer, to build more complex GNN architectures.

To use the PyTorch Geometric implementation of GraphSAGE, we first need to import the necessary libraries and load the dataset. In this example, we will use the Cora dataset, which is a popular benchmark dataset for node classification tasks.

import torch
from torch_geometric.datasets import Planetoid
from torch_geometric.data import Data
from torch_geometric.nn import GraphSAGE

# Load the Cora dataset
dataset = Planetoid(root='./cora', name='Cora')

Next, we need to define the GraphSAGE model architecture. In this example, we will use a simple GraphSAGE model with two layers, each with 64 hidden units.

# Define the GraphSAGE model architecture
class GraphSAGEModel(torch.nn.Module):
    def __init__(self):
        super(GraphSAGEModel, self).__init__()
        self.sage1 = GraphSAGE(dataset.num_features, 64)
        self.sage2 = GraphSAGE(64, dataset.num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = self.sage1(x, edge_index)
        x = torch.relu(x)
        x = self.sage2(x, edge_index)

        return torch.log_softmax(x, dim=1)

# Initialize the GraphSAGE model
model = GraphSAGEModel()

Once we have defined the GraphSAGE model architecture, we can train the model using the PyTorch Geometric train function. In this example, we will use the Adam optimizer and cross-entropy loss.

# Define the training loop
def train(model, device, data, optimizer, epoch):
    model.train()
    optimizer.zero_grad()
    out = model(data)
    loss = torch.nn.NLLLoss()(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()

# Train the GraphSAGE model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
data = dataset[0].to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

for epoch in range(200):
    loss = train(model, device, data, optimizer, epoch)
    print(f'Epoch: {epoch+1}, Loss: {loss:.4f}')

Evaluation of GraphSAGE Models

Once we have trained the GraphSAGE model, we can evaluate its performance using the PyTorch Geometric evaluate function. In this example, we will use the accuracy metric to evaluate the model’s performance.

# Define the evaluation function
def evaluate(model, device, data):
    model.eval()
    _, pred = model(data).max(dim=1)
    correct = int(pred[data.test_mask].eq(data.y[data.test_mask]).sum().item())
    return correct / int(data.test_mask.sum())

# Evaluate the GraphSAGE model
test_acc = evaluate(model, device, data)
print(f'Test Accuracy: {test_acc:.4f}')

Comparison with Other GNN Architectures

GraphSAGE is just one of many GNN architectures that can be used for node classification tasks. Other popular GNN architectures include GCN, GAT, and Graph Attention Network. In this section, we will compare the performance of GraphSAGE with these other GNN architectures on the Cora dataset.

GCN Architecture

GCN is a popular GNN architecture that uses a graph convolutional layer to learn node representations. The GCN architecture can be implemented using the PyTorch Geometric GCNConv layer.

# Define the GCN model architecture
class GCNModel(torch.nn.Module):
    def __init__(self):
        super(GCNModel, self).__init__()
        self.conv1 = GCNConv(dataset.num_features, 64)
        self.conv2 = GCNConv(64, dataset.num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = self.conv1(x, edge_index)
        x = torch.relu(x)
        x = self.conv2(x, edge_index)

        return torch.log_softmax(x, dim=1)

# Initialize the GCN model
model = GCNModel()

GAT Architecture

GAT is a popular GNN architecture that uses a graph attention layer to learn node representations. The GAT architecture can be implemented using the PyTorch Geometric GATConv layer.

# Define the GAT model architecture
class GATModel(torch.nn.Module):
    def __init__(self):
        super(GATModel, self).__init__()
        self.conv1 = GATConv(dataset.num_features, 64)
        self.conv2 = GATConv(64, dataset.num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = self.conv1(x, edge_index)
        x = torch.relu(x)
        x = self.conv2(x, edge_index)

        return torch.log_softmax(x, dim=1)

# Initialize the GAT model
model = GATModel()