Exploring Graph Attention Networks (GATs) in PyTorch for Node Classification

Graph Attention Networks (GATs) have emerged as a potent tool in the realm of graph machine learning. By leveraging the mechanism of attention, GATs can dynamically focus on the most pertinent parts of a graph, which proves advantageous in tasks like node classification. This article delves into the implementation of GATs using PyTorch, elucidating how they can be set up for node classification tasks.

Introduction to Graph Attention Networks

Traditional graph neural networks (GNNs) like Graph Convolutional Networks (GCNs) rely on fixed weights of neighboring nodes. However, GATs introduce a mechanism to compute attention scores, allowing the model to individually weigh neighbor contributions based on their relevance to the node being processed. This adaptability makes GATs highly effective for diverse graph-based tasks.

PyTorch and PyTorch Geometric

For graph network implementations, PyTorch Geometric is a well-liked library that sits atop PyTorch. It simplifies handling of complex graph data structures and provides an extensive suite of predefined layers and utilities that make writing GATs straightforward.

Before we jump into code, ensure that PyTorch and PyTorch Geometric are installed:

pip install torch torchvision torchaudio
pip install torch-geometric

Implementing GAT with PyTorch Geometric

We will begin by creating a simple Graph Attention Network for a node classification task. Let us assume we use the Cora dataset, a citation network dataset, commonly utilized in node classification experiments. We'll leverage PyTorch Geometric's built-in dataset functionality.

import torch
import torch.nn.functional as F
from torch_geometric.nn import GATConv
from torch_geometric.datasets import Planetoid
from torch_geometric.data import DataLoader

# Load Cora dataset
dataset = Planetoid(root='/tmp/Cora', name='Cora')

Now that our dataset is ready, let's define the GAT model. Our GAT model will have two GAT layers. We will also use dropout to mitigate overfitting:

class GAT(torch.nn.Module):
    def __init__(self, data):
        super(GAT, self).__init__()
        self.gat1 = GATConv(dataset.num_features, 8, heads=8, dropout=0.6)
        self.gat2 = GATConv(8 * 8, dataset.num_classes, heads=1, concat=True, dropout=0.6)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = F.dropout(x, p=0.6, training=self.training)
        x = F.elu(self.gat1(x, edge_index))
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.gat2(x, edge_index)
        return F.log_softmax(x, dim=1)

Training the Network

With the model defined, the next step is to train it. Define a training loop that performs backpropagation on the network:

model = GAT(dataset)
optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

def train():
    model.train()
    optimizer.zero_grad()
    out = model(dataset[0])
    loss = F.nll_loss(out[dataset[0].train_mask], dataset[0].y[dataset[0].train_mask])
    loss.backward()
    optimizer.step()
    return loss

for epoch in range(200):
    loss = train()
    print('Epoch {:03d}, Loss: {:.4f}'.format(epoch, loss.item()))

The above code defines a simple training loop. The training process continues for 200 epochs or until the learning convergence is satisfactory.

Evaluating the Model

To evaluate our trained model, a simple evaluating function can be written as:

def test():
    model.eval()
    out = model(dataset[0])
    pred = out.argmax(dim=1)
    test_correct = pred[dataset[0].test_mask] == dataset[0].y[dataset[0].test_mask]
    test_acc = int(test_correct.sum()) / int(dataset[0].test_mask.sum())
    return test_acc

acc = test()
print('Accuracy: {:.4f}'.format(acc))

We calculate the accuracy by comparing the predicted node labels against the true labels for the nodes within the test mask.

Conclusion

Graph Attention Networks provide robust capabilities for tackling complex graph-related tasks, particularly node classification. By employing a multi-head attention mechanism, GATs allow the model to focus on different parts of its neighbors, often leading to improved performance over models like GCNs. PyTorch Geometric streamlines the process of building such complex models, making the approach practical and effective.