Building Your First Graph Convolutional Network (GCN) with PyTorch

Graph Convolutional Networks (GCNs) have become a prominent method for machine learning on graph-structured data. PyTorch, with its dynamic computation graph and simple API, is an excellent choice for implementing GCNs. In this tutorial, we will guide you through building your first GCN using PyTorch.

Prerequisites

Before we start, make sure you have PyTorch installed on your system. You can install it using:

pip install torch torchvision

Additionally, we will use some basic Python libraries such as numpy for numerical computations:

pip install numpy

Step 1: Understanding the Graph Structure

Graphs consist of nodes and edges connecting these nodes. Representing these in code, we can start by simulating a simple graph using adjacency matrices and node features.

Suppose we have a graph with three nodes:

import numpy as np

# Example adjacency matrix
adjacency_matrix = np.array([[0, 1, 1],
                            [1, 0, 1],
                            [1, 1, 0]])

# Node features (e.g., 2 features per node)
features = np.array([[1, 2],
                     [2, 3],
                     [3, 4]])

Step 2: Designing the GCN Layer

To perform graph convolution, we need to define a GCN layer. The paper “Semi-Supervised Classification with Graph Convolutional Networks” introduces a simple GCN layer defined as:

$$H^{(l+1)} = \sigma(\tilde{A}H^{(l)}W^{(l)})$$

where \(\tilde{A}\) is the normalized adjacency matrix, \(H^{(l)}\) is the input feature matrix, \(W^{(l)}\) is a trainable weight matrix, and \(\sigma\) is an activation function (e.g., ReLU).

Here is a simple implementation in PyTorch:

import torch
import torch.nn as nn
import torch.nn.functional as F

class GCNLayer(nn.Module):
    def __init__(self, in_features, out_features):
        super(GCNLayer, self).__init__()
        self.weight = nn.Parameter(torch.FloatTensor(in_features, out_features))
        nn.init.xavier_uniform_(self.weight.data, gain=1.0)

    def forward(self, adj, h):
        # Apply graph convolution
        h = torch.mm(adj, h)
        h = torch.mm(h, self.weight)
        return F.relu(h)

Step 3: Building the GCN Model

We can now stack multiple GCN layers to build a simple GCN model. Here, we create a two-layer GCN for node classification:

class GCNModel(nn.Module):
    def __init__(self, n_features, n_classes):
        super(GCNModel, self).__init__()
        self.gcn1 = GCNLayer(n_features, 16)  # Hidden layer
        self.gcn2 = GCNLayer(16, n_classes)   # Output layer

    def forward(self, adj, features):
        x = self.gcn1(adj, features)
        x = self.gcn2(adj, x)
        return F.log_softmax(x, dim=1)

Step 4: Training the Model

With the model ready, you can conduct forward and backward propagation for training:

# Initialized model, loss, and optimizer
model = GCNModel(n_features=2, n_classes=3)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()

# Convert data to torch tensors
adj = torch.FloatTensor(adjacency_matrix)
features = torch.FloatTensor(features)
labels = torch.LongTensor([0, 1, 2])  # Assuming labels for the nodes

# Training loop
model.train()
for epoch in range(200):
    optimizer.zero_grad()
    output = model(adj, features)
    loss = loss_fn(output, labels)
    loss.backward()
    optimizer.step()

    if epoch % 20 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item()}')

Conclusion

Congratulations! You've just built your first Graph Convolutional Network using PyTorch. While we implemented a simple model here, GCNs are a versatile class of architectures applicable to numerous graph-based tasks such as node classification, link prediction, and graph-level classification.