In Chapters 1 and 2, you learned the fundamentals of graph data and how to perform node classification with a Graph Convolutional Network (GCN). In this chapter, we go further into more expressive Graph Neural Network (GNN) architectures. We'll systematically work through Graph Attention Networks (GAT), which learn the importance of relationships between nodes; GraphSAGE, which scales to large graphs; graph pooling, which aggregates an entire graph into a single vector; and application tasks such as graph classification, link prediction, and heterogeneous graphs — all with working code.
Learning Objectives
- ✅ Explain how Graph Attention Networks (GAT) work
- ✅ Understand GraphSAGE's sampling strategy
- ✅ Obtain a whole-graph representation using graph pooling
- ✅ Implement a graph classification task
- ✅ Understand the basics of link prediction tasks
- ✅ Know how to work with heterogeneous graphs
1. Graph Attention Networks (GAT)
Graph Attention Network (GAT) is a GNN architecture that, instead of averaging all neighboring nodes with equal weight as GCN does, learns a different weight of importance for each node using an attention mechanism. It's easiest to understand as an application to graph structures of the same idea behind Self-Attention, used in Transformers for natural language processing.
Basic Principles of GAT
In GAT, the attention coefficient \(\alpha_{ij}\) between node \(i\) and its neighboring node \(j\) is learned from the feature vectors of both nodes.
First, an attention score \(e_{ij}\) is computed as follows:
$$e_{ij} = \text{LeakyReLU}\left(\mathbf{a}^T [\mathbf{W}\mathbf{h}_i \, \| \, \mathbf{W}\mathbf{h}_j]\right)$$
Here \(\mathbf{W}\) is a learnable weight matrix, \(\mathbf{a}\) is a learnable attention vector, and \(\|\) denotes vector concatenation. Next, we normalize across all neighboring nodes by applying the softmax function:
$$\alpha_{ij} = \frac{\exp(e_{ij})}{\sum_{k \in \mathcal{N}(i)} \exp(e_{ik})}$$
Finally, we compute the weighted sum of neighbor feature vectors, weighted by the attention coefficients, to obtain the new feature vector for node \(i\):
$$\mathbf{h}_i' = \sigma\left(\sum_{j \in \mathcal{N}(i)} \alpha_{ij} \mathbf{W} \mathbf{h}_j\right)$$
💡 GCN vs. GAT
GCN aggregates neighbor features using a fixed weight \(1/\sqrt{d_i d_j}\) determined by node degree (the number of neighbors), whereas GAT learns importance directly from the data. This makes it possible to suppress the influence of noisy neighboring nodes or emphasize important ones, giving GAT greater representational power than GCN on many tasks.
Multi-Head Attention
Like the Transformer, GAT uses multi-head attention: it learns several independent attention mechanisms in parallel and concatenates (or averages) their results, which stabilizes training and increases representational power. In torch_geometric.nn.GATConv, you specify the number of heads with the heads argument.
Implementing GATConv
Let's implement node classification with GATConv using the Cora dataset (a citation network).
import torch
import torch.nn.functional as F
from torch_geometric.nn import GATConv
from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]
class GAT(torch.nn.Module):
def __init__(self, num_features, num_classes, hidden_channels=8, heads=8):
super().__init__()
# Layer 1: use 8-head attention and concatenate the outputs (concat=True)
self.conv1 = GATConv(num_features, hidden_channels, heads=heads, dropout=0.6)
# Layer 2: use a single head so the output dimension matches the number of classes
self.conv2 = GATConv(hidden_channels * heads, num_classes, heads=1,
concat=False, dropout=0.6)
def forward(self, x, edge_index):
x = F.dropout(x, p=0.6, training=self.training)
x = F.elu(self.conv1(x, edge_index))
x = F.dropout(x, p=0.6, training=self.training)
x = self.conv2(x, edge_index)
return F.log_softmax(x, dim=1)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GAT(dataset.num_features, dataset.num_classes).to(device)
data = data.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)
model.train()
for epoch in range(100):
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()
if (epoch + 1) % 20 == 0:
print(f'Epoch {epoch+1:03d}, Loss: {loss:.4f}')
model.eval()
with torch.no_grad():
pred = model(data.x, data.edge_index).argmax(dim=1)
test_acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean().item()
print(f'Test Accuracy: {test_acc:.4f}')
Sample Output:
Epoch 020, Loss: 1.2345
Epoch 040, Loss: 0.7823
Epoch 060, Loss: 0.5412
Epoch 080, Loss: 0.4501
Epoch 100, Loss: 0.3987
Test Accuracy: 0.8280
💡 Visualizing Attention Weights
If you pass return_attention_weights=True when calling GATConv's forward, you can retrieve the attention coefficient \(\alpha_{ij}\) assigned to each edge. This lets you analyze which nodes the model relied on most for its predictions, improving the model's interpretability.
2. GraphSAGE (Sampling-Based GNN)
GraphSAGE (Graph SAmple and aggreGatE) is a GNN designed to scale to graphs with millions or even billions of nodes. Whereas GCN and GAT assume full-batch learning, processing the entire graph at once using every neighboring node, GraphSAGE samples a fixed number of neighbors per node and trains in mini-batches, keeping memory usage and computational cost under control.
The GraphSAGE Algorithm
Each layer of GraphSAGE performs the following two steps.
- Sampling: for each node, randomly draw a fixed number (e.g., 10) of neighboring nodes
- Aggregation: aggregate the features of the sampled neighbors and combine them with the node's own features to produce an updated representation
$$\mathbf{h}_i^{(k)} = \sigma\left(\mathbf{W}^{(k)} \cdot \text{CONCAT}\left(\mathbf{h}_i^{(k-1)}, \; \text{AGG}_k\left(\{\mathbf{h}_j^{(k-1)} : j \in \mathcal{N}_s(i)\}\right)\right)\right)$$
Here \(\mathcal{N}_s(i)\) is the set of sampled neighbors of node \(i\), and \(\text{AGG}_k\) is the aggregation function.
| Aggregation Function | Description |
|---|---|
| Mean Aggregator | Takes the average of neighbor feature vectors. The simplest and least expensive option |
| Pooling Aggregator | Passes each neighbor feature vector through an MLP, then takes the elementwise maximum |
| LSTM Aggregator | Feeds neighbors into an LSTM in a random order. Highly expressive, but has an order-dependency issue |
💡 Why Sampling Works
If you process an entire graph in a single forward pass, as GCN does, the set of nodes you need to reference expands exponentially with every hop you follow (a problem known as neighbor explosion). By fixing the number of samples at each layer, GraphSAGE keeps this explosion in check and enables mini-batch training with memory usage that doesn't depend on the size of the overall graph.
Implementing a SAGEConv Model
import torch
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv
from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]
class GraphSAGE(torch.nn.Module):
def __init__(self, num_features, num_classes, hidden_channels=16):
super().__init__()
self.conv1 = SAGEConv(num_features, hidden_channels)
self.conv2 = SAGEConv(hidden_channels, num_classes)
def forward(self, x, edge_index):
x = F.relu(self.conv1(x, edge_index))
x = F.dropout(x, p=0.5, training=self.training)
x = self.conv2(x, edge_index)
return x
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GraphSAGE(dataset.num_features, dataset.num_classes).to(device)
print(model)
Mini-Batch Training with NeighborLoader
To simulate training on a large graph, we use torch_geometric.loader.NeighborLoader to sample neighboring nodes and generate mini-batches. num_neighbors=[10, 5] means "sample up to 10 neighbors at the first layer, and up to 5 neighbors at the second layer."
from torch_geometric.loader import NeighborLoader
# Generate mini-batches starting from the training nodes, sampling neighbors up to 2 hops away
train_loader = NeighborLoader(
data,
num_neighbors=[10, 5],
batch_size=128,
input_nodes=data.train_mask,
shuffle=True,
)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
model.train()
for epoch in range(20):
total_loss = 0
for batch in train_loader:
batch = batch.to(device)
optimizer.zero_grad()
out = model(batch.x, batch.edge_index)
# The first batch_size entries in the batch are the "seed nodes";
# the rest were added by sampling and are not used when computing the loss
loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])
loss.backward()
optimizer.step()
total_loss += loss.item()
if (epoch + 1) % 5 == 0:
print(f'Epoch {epoch+1:03d}, Loss: {total_loss / len(train_loader):.4f}')
Evaluation
model.eval()
data_device = data.to(device)
with torch.no_grad():
out = model(data_device.x, data_device.edge_index)
pred = out.argmax(dim=1)
test_acc = (pred[data_device.test_mask] == data_device.y[data_device.test_mask]).float().mean().item()
print(f'Test Accuracy: {test_acc:.4f}')
Sample Output:
Epoch 005, Loss: 1.5432
Epoch 010, Loss: 0.9821
Epoch 015, Loss: 0.6543
Epoch 020, Loss: 0.4987
Test Accuracy: 0.7920
⚠️ Impact of Sampling on Accuracy
Because mini-batch training doesn't use the entire graph's information at once, accuracy can be slightly lower than with full-batch GCN or GAT. That said, for graphs with millions of nodes, full-batch training is often simply impossible due to memory constraints, making GraphSAGE's sampling strategy indispensable in such settings.
3. Graph Pooling Methods
The GCN, GAT, and GraphSAGE models we've seen so far all operate at the node level, updating each node's feature vector. But for tasks like predicting the toxicity of an entire molecule or classifying a protein's function, where we want a single prediction for the "whole graph," we need to aggregate all the per-node features into one vector representing the entire graph. This aggregation step is called graph pooling or readout.
Global Pooling
The most basic form of pooling is global pooling, which aggregates the features of all nodes using a simple statistic.
| Function | Aggregation Method | Characteristics |
|---|---|---|
global_mean_pool |
Average of all node features | Robust to graph size; the most widely used option |
global_max_pool |
Elementwise maximum across all node features | Emphasizes salient features; sensitive to outliers |
global_add_pool |
Sum of all node features | Preserves information about the number of nodes (graph size) |
import torch
from torch_geometric.nn import global_mean_pool, global_max_pool, global_add_pool
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader
dataset = TUDataset(root='/tmp/MUTAG', name='MUTAG')
loader = DataLoader(dataset, batch_size=4, shuffle=True)
batch = next(iter(loader))
print(f'Batch: {batch.num_graphs} graphs, {batch.num_nodes} nodes total')
# Treat node features directly as the graph representation and compare the output shape of each pooling method
mean_pooled = global_mean_pool(batch.x, batch.batch)
max_pooled = global_max_pool(batch.x, batch.batch)
sum_pooled = global_add_pool(batch.x, batch.batch)
print(f'Mean pooling shape: {mean_pooled.shape}')
print(f'Max pooling shape: {max_pooled.shape}')
print(f'Sum pooling shape: {sum_pooled.shape}')
Sample Output:
Batch: 4 graphs, 71 nodes total
Mean pooling shape: torch.Size([4, 7])
Max pooling shape: torch.Size([4, 7])
Sum pooling shape: torch.Size([4, 7])
Notice that the output shape is [num_graphs, feature_dim]. Using the batch attribute (an index recording which graph each node belongs to), several graphs with different numbers of nodes were combined into a single batch, yet each one is aggregated correctly and separately.
Hierarchical Pooling
Global pooling aggregates all nodes at once, whereas hierarchical pooling progressively reduces the number of nodes between GNN layers, much like the pooling layers of a CNN. TopKPooling is a representative technique that scores each node's importance with a learnable scoring function and keeps only the top \(k\) nodes (specified via a ratio).
import torch
import torch.nn.functional as F
from torch_geometric.nn import GraphConv, TopKPooling, global_mean_pool
class HierarchicalPoolNet(torch.nn.Module):
def __init__(self, num_features, num_classes, hidden_channels=64):
super().__init__()
self.conv1 = GraphConv(num_features, hidden_channels)
# Keep only the top 80% of nodes by importance score
self.pool1 = TopKPooling(hidden_channels, ratio=0.8)
self.conv2 = GraphConv(hidden_channels, hidden_channels)
self.lin = torch.nn.Linear(hidden_channels, num_classes)
def forward(self, x, edge_index, batch):
x = F.relu(self.conv1(x, edge_index))
# Reduce the number of nodes via pooling while updating edge_index and batch to stay consistent
x, edge_index, _, batch, _, _ = self.pool1(x, edge_index, None, batch)
x = F.relu(self.conv2(x, edge_index))
x = global_mean_pool(x, batch)
return self.lin(x)
model = HierarchicalPoolNet(num_features=7, num_classes=2)
print(model)
💡 Which Should You Use?
Global pooling is simple to implement and computationally cheap, making it the right first choice to try. Hierarchical pooling has greater representational power but adds model complexity and can make training less stable, so it's worth considering when global pooling doesn't give sufficient accuracy, or when you want to reduce computation on large graphs.
4. Graph Classification Task
Graph classification is the task of predicting a single label for an entire graph. Typical examples include predicting a molecule's toxicity or activity and classifying a protein's function. Here we implement a graph classification model that combines GATConv and SAGEConv from start to finish, using the MUTAG dataset (molecular graphs of mutagenic compounds, one of the TUDataset collections).
Preparing the Dataset
import torch
import torch.nn.functional as F
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader
from torch_geometric.nn import GATConv, SAGEConv, global_mean_pool
torch.manual_seed(42)
# MUTAG: 188 molecular graphs, binary classification (mutagenic or not)
dataset = TUDataset(root='/tmp/MUTAG', name='MUTAG').shuffle()
train_dataset = dataset[:150]
test_dataset = dataset[150:]
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32)
print(f'Number of training graphs: {len(train_dataset)}')
print(f'Number of test graphs: {len(test_dataset)}')
print(f'Number of node features: {dataset.num_node_features}')
print(f'Number of classes: {dataset.num_classes}')
Sample Output:
Number of training graphs: 150
Number of test graphs: 38
Number of node features: 7
Number of classes: 2
Defining the Model
This architecture uses GATConv to learn the importance of relationships between nodes, then SAGEConv to aggregate neighborhood information, and finally global_mean_pool to obtain a vector representation of the entire graph.
class GraphClassifier(torch.nn.Module):
def __init__(self, num_features, num_classes, hidden_channels=64):
super().__init__()
# GAT layer: extract features that account for node-to-node importance via multi-head attention
self.gat = GATConv(num_features, hidden_channels, heads=4, concat=True, dropout=0.2)
# GraphSAGE layer: aggregate neighborhood information to refine the features
self.sage = SAGEConv(hidden_channels * 4, hidden_channels)
self.lin = torch.nn.Linear(hidden_channels, num_classes)
def forward(self, x, edge_index, batch):
x = F.elu(self.gat(x, edge_index))
x = F.dropout(x, p=0.2, training=self.training)
x = F.relu(self.sage(x, edge_index))
# Aggregate the whole graph into a single vector (readout)
x = global_mean_pool(x, batch)
x = F.dropout(x, p=0.5, training=self.training)
return self.lin(x)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GraphClassifier(dataset.num_node_features, dataset.num_classes).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
print(model)
Training and Evaluation
def train():
model.train()
total_loss = 0
for data in train_loader:
data = data.to(device)
optimizer.zero_grad()
out = model(data.x, data.edge_index, data.batch)
loss = F.cross_entropy(out, data.y)
loss.backward()
optimizer.step()
total_loss += loss.item() * data.num_graphs
return total_loss / len(train_loader.dataset)
@torch.no_grad()
def test(loader):
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)
for epoch in range(1, 101):
loss = train()
if epoch % 20 == 0:
train_acc = test(train_loader)
test_acc = test(test_loader)
print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, '
f'Train Acc: {train_acc:.4f}, Test Acc: {test_acc:.4f}')
Sample Output:
Epoch 020, Loss: 0.5893, Train Acc: 0.7133, Test Acc: 0.6842
Epoch 040, Loss: 0.4721, Train Acc: 0.7867, Test Acc: 0.7368
Epoch 060, Loss: 0.4102, Train Acc: 0.8200, Test Acc: 0.7632
Epoch 080, Loss: 0.3654, Train Acc: 0.8467, Test Acc: 0.7895
Epoch 100, Loss: 0.3389, Train Acc: 0.8600, Test Acc: 0.8158
🎉 Graph Classification Model Complete
We achieved roughly 80% test accuracy on the MUTAG dataset. The pattern of using a GAT layer to weight important atoms (nodes), a SAGEConv layer to aggregate local structure, and a final global pooling step to obtain a whole-molecule representation is a typical pipeline for molecular property prediction tasks.
5. Link Prediction Task
Link prediction is the task of predicting whether an edge (link) exists between two nodes in a graph. It has wide-ranging applications, including friend recommendations in social networks, product recommendations in recommender systems, and completing missing relations in knowledge graphs.
Task Design
In link prediction, existing edges are sampled as positive samples and non-existent node pairs as negative samples, and the problem is trained as binary classification. When splitting edges into training, validation, and test sets, the split of the graph structure and the split of labels must be done together to avoid information leakage. PyTorch Geometric's RandomLinkSplit automates this for us.
import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.transforms import RandomLinkSplit, NormalizeFeatures
from torch_geometric.nn import GCNConv
from torch_geometric.utils import negative_sampling
dataset = Planetoid(root='/tmp/Cora', name='Cora', transform=NormalizeFeatures())
data = dataset[0]
# Split edges into train/val/test sets (negative edges are also generated automatically)
transform = RandomLinkSplit(
num_val=0.05,
num_test=0.1,
is_undirected=True,
add_negative_train_samples=False,
)
train_data, val_data, test_data = transform(data)
print(f'Train message-passing edges: {train_data.edge_index.shape}')
print(f'Train supervision edges: {train_data.edge_label_index.shape}')
print(f'Val supervision edges: {val_data.edge_label_index.shape}')
print(f'Test supervision edges: {test_data.edge_label_index.shape}')
💡 Message-Passing Edges vs. Supervision Edges
RandomLinkSplit separately manages edge_index, which the GNN uses to propagate features (message passing), and edge_label_index, which is used for computing loss and evaluation (supervision). This prevents the leak in which the model would otherwise get to see the very edge it's supposed to predict.
Implementing an Encoder-Decoder Model
We build a model consisting of an encoder that computes node embeddings with GCNConv, and a decoder that computes the probability of an edge's existence from the dot product of two node embeddings.
class LinkPredGCN(torch.nn.Module):
def __init__(self, num_features, hidden_channels=128, out_channels=64):
super().__init__()
self.conv1 = GCNConv(num_features, hidden_channels)
self.conv2 = GCNConv(hidden_channels, out_channels)
def encode(self, x, edge_index):
x = F.relu(self.conv1(x, edge_index))
return self.conv2(x, edge_index)
def decode(self, z, edge_label_index):
# Use the dot product of the two node embeddings as the score for edge existence
src, dst = edge_label_index
return (z[src] * z[dst]).sum(dim=-1)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = LinkPredGCN(dataset.num_features).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
train_data = train_data.to(device)
val_data = val_data.to(device)
test_data = test_data.to(device)
Training and Evaluation
During training, we resample new negative edges every epoch so the model doesn't overfit to a fixed set of negatives. For the evaluation metric we use ROC-AUC (Receiver Operating Characteristic - Area Under the Curve), a common choice for binary classification.
from sklearn.metrics import roc_auc_score
def train():
model.train()
optimizer.zero_grad()
z = model.encode(train_data.x, train_data.edge_index)
# Randomly resample negative edges (non-existent edges) every epoch
neg_edge_index = negative_sampling(
edge_index=train_data.edge_index,
num_nodes=train_data.num_nodes,
num_neg_samples=train_data.edge_label_index.size(1),
method='sparse',
)
edge_label_index = torch.cat([train_data.edge_label_index, neg_edge_index], dim=1)
edge_label = torch.cat([
train_data.edge_label,
train_data.edge_label.new_zeros(neg_edge_index.size(1)),
], dim=0)
out = model.decode(z, edge_label_index).view(-1)
loss = F.binary_cross_entropy_with_logits(out, edge_label)
loss.backward()
optimizer.step()
return loss.item()
@torch.no_grad()
def test(data):
model.eval()
z = model.encode(data.x, data.edge_index)
out = model.decode(z, data.edge_label_index).view(-1).sigmoid()
return roc_auc_score(data.edge_label.cpu().numpy(), out.cpu().numpy())
for epoch in range(1, 101):
loss = train()
if epoch % 20 == 0:
val_auc = test(val_data)
test_auc = test(test_data)
print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, '
f'Val AUC: {val_auc:.4f}, Test AUC: {test_auc:.4f}')
Sample Output:
Epoch 020, Loss: 0.6123, Val AUC: 0.7845, Test AUC: 0.7791
Epoch 040, Loss: 0.5234, Val AUC: 0.8412, Test AUC: 0.8356
Epoch 060, Loss: 0.4756, Val AUC: 0.8723, Test AUC: 0.8681
Epoch 080, Loss: 0.4489, Val AUC: 0.8891, Test AUC: 0.8834
Epoch 100, Loss: 0.4321, Val AUC: 0.8956, Test AUC: 0.8902
🎉 Link Prediction Model Complete
We achieved a test AUC of roughly 0.89 on Cora's citation network. Since an AUC of 0.5 is equivalent to random guessing and 1.0 means perfect prediction, a score of 0.89 shows that the model can distinguish "paper pairs that are actually cited" from "pairs that are not" quite accurately.
6. Working with Heterogeneous Graphs
All the graphs we've worked with so far consisted of a single type of node (e.g., paper nodes only) — a homogeneous graph. In the real world, however, many graphs are heterogeneous graphs that contain multiple types of nodes and edges, such as users and products, authors and papers, or atoms and bond types.
The HeteroData Object
PyTorch Geometric represents heterogeneous graphs with a HeteroData object. Node features are stored per node type, and edges are managed as triples of (source node type, relation name, target node type). Let's build a bipartite graph as an example, using a recommendation system in which users rate movies.
import torch
from torch_geometric.data import HeteroData
data = HeteroData()
# Node features (the dimensionality can differ across node types)
data['user'].x = torch.randn(5, 16) # 5 users, 16-dimensional features
data['movie'].x = torch.randn(8, 32) # 8 movies, 32-dimensional features
# Edges: defined as (source node type, relation name, target node type) triples
data['user', 'rates', 'movie'].edge_index = torch.tensor([
[0, 0, 1, 2, 3, 4],
[0, 1, 1, 2, 5, 7],
], dtype=torch.long)
# Explicitly add the reverse edges too (message passing is treated as directed)
data['movie', 'rated_by', 'user'].edge_index = \
data['user', 'rates', 'movie'].edge_index.flip(0)
print(data)
print(f'Node types: {data.node_types}')
print(f'Edge types: {data.edge_types}')
Sample Output:
HeteroData(
user={ x=[5, 16] },
movie={ x=[8, 32] },
(user, rates, movie)={ edge_index=[2, 6] },
(movie, rated_by, user)={ edge_index=[2, 6] }
)
Node types: ['user', 'movie']
Edge types: [('user', 'rates', 'movie'), ('movie', 'rated_by', 'user')]
Heterogeneous GNNs with HeteroConv
Using torch_geometric.nn.HeteroConv, you can concisely build a model that applies a separate GNN layer (here, SAGEConv) to each edge type and aggregates the results for each node type. When the input dimension differs across node types, specifying -1 as in SAGEConv((-1, -1), hidden_channels) lets the input dimension be inferred automatically on the first forward pass.
from torch_geometric.nn import HeteroConv, SAGEConv, Linear
class HeteroGNN(torch.nn.Module):
def __init__(self, hidden_channels, out_channels):
super().__init__()
# Apply a separate SAGEConv for each edge type, then aggregate the results per node type
self.conv1 = HeteroConv({
('user', 'rates', 'movie'): SAGEConv((-1, -1), hidden_channels),
('movie', 'rated_by', 'user'): SAGEConv((-1, -1), hidden_channels),
}, aggr='sum')
self.lin_user = Linear(hidden_channels, out_channels)
self.lin_movie = Linear(hidden_channels, out_channels)
def forward(self, x_dict, edge_index_dict):
x_dict = self.conv1(x_dict, edge_index_dict)
x_dict = {key: x.relu() for key, x in x_dict.items()}
return {
'user': self.lin_user(x_dict['user']),
'movie': self.lin_movie(x_dict['movie']),
}
model = HeteroGNN(hidden_channels=16, out_channels=4)
out_dict = model(data.x_dict, data.edge_index_dict)
print(f"User embeddings shape: {out_dict['user'].shape}")
print(f"Movie embeddings shape: {out_dict['movie'].shape}")
Sample Output:
User embeddings shape: torch.Size([5, 4])
Movie embeddings shape: torch.Size([8, 4])
💡 The Handy to_hetero() Conversion Function
If you already have a GNN model built for homogeneous graphs (e.g., GCN, GAT), the torch_geometric.nn.to_hetero() function can automatically convert it into a heterogeneous-graph-compatible version without rewriting the model definition. It's worth knowing both approaches — building from scratch with HeteroConv, and converting an existing model with to_hetero() — so you can pick whichever fits the situation.
⚠️ A Caution on Heterogeneous Graph Design
If you forget to explicitly add the reverse edge type (e.g., rated_by) in a heterogeneous graph, information will only propagate in one direction. In practice, torch_geometric.transforms.ToUndirected() is often used because it automatically adds reverse edges for every edge type.
Exercises
Exercise 1: Compare Different Numbers of GAT Heads
Using the GAT model built in "Implementing GATConv," train on the Cora dataset with heads set to 2, 4, and 8, and compare the resulting test accuracy. Discuss how the number of heads affects accuracy and training time.
Show Answer
Implementation notes:
for heads in [2, 4, 8]:
model = GAT(dataset.num_features, dataset.num_classes, heads=heads).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)
model.train()
for epoch in range(100):
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()
model.eval()
with torch.no_grad():
pred = model(data.x, data.edge_index).argmax(dim=1)
acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean().item()
print(f'heads={heads}: Test Accuracy = {acc:.4f}')
Discussion: Increasing the number of heads generally lets the model learn several distinct "attention patterns" in parallel, improving representational power and making accuracy more stable. On the other hand, the parameter count and computational cost grow linearly with the number of heads, so training time per epoch increases as well. On a small dataset like Cora, accuracy typically plateaus around 8 heads, and adding more heads beyond that tends not to yield accuracy gains that justify the extra compute.
Exercise 2: Add Multiple Kinds of Global Pooling to the Graph Classifier
Modify the GraphClassifier implemented in "Graph Classification Task" so that it concatenates the outputs of global_mean_pool and global_max_pool before feeding them into the fully connected layer. Check whether this improves accuracy.
Show Answer
Example implementation:
from torch_geometric.nn import global_mean_pool, global_max_pool
class GraphClassifierV2(torch.nn.Module):
def __init__(self, num_features, num_classes, hidden_channels=64):
super().__init__()
self.gat = GATConv(num_features, hidden_channels, heads=4, concat=True, dropout=0.2)
self.sage = SAGEConv(hidden_channels * 4, hidden_channels)
# Concatenating mean and max doubles the dimensionality
self.lin = torch.nn.Linear(hidden_channels * 2, num_classes)
def forward(self, x, edge_index, batch):
x = F.elu(self.gat(x, edge_index))
x = F.dropout(x, p=0.2, training=self.training)
x = F.relu(self.sage(x, edge_index))
x_mean = global_mean_pool(x, batch)
x_max = global_max_pool(x, batch)
x = torch.cat([x_mean, x_max], dim=1)
x = F.dropout(x, p=0.5, training=self.training)
return self.lin(x)
Discussion: global_mean_pool captures the "average characteristics" of the whole graph, while global_max_pool captures its "most salient characteristics," so concatenating both retains more information in the graph representation. On a small dataset like MUTAG the improvement may only be a few percentage points, but on datasets with more complex molecular structures the effect tends to be more pronounced.
Exercise 3: Add a New Node Type to the Heterogeneous Graph
Extend the example in "The HeteroData Object" by adding a new node type 'genre' (3 movie genres, with 8-dimensional features), and connect movies to genres with an edge type ('movie', 'has_genre', 'genre'). Then add a SAGEConv for this new edge type to the HeteroGNN model's HeteroConv.
Show Answer
Example implementation:
import torch
from torch_geometric.data import HeteroData
from torch_geometric.nn import HeteroConv, SAGEConv, Linear
data = HeteroData()
data['user'].x = torch.randn(5, 16)
data['movie'].x = torch.randn(8, 32)
data['genre'].x = torch.randn(3, 8) # New node type: genre
data['user', 'rates', 'movie'].edge_index = torch.tensor([
[0, 0, 1, 2, 3, 4],
[0, 1, 1, 2, 5, 7],
], dtype=torch.long)
data['movie', 'rated_by', 'user'].edge_index = \
data['user', 'rates', 'movie'].edge_index.flip(0)
# Connect each movie (0-7) to its corresponding genre (0-2)
data['movie', 'has_genre', 'genre'].edge_index = torch.tensor([
[0, 1, 2, 3, 4, 5, 6, 7],
[0, 0, 1, 1, 2, 2, 0, 1],
], dtype=torch.long)
class HeteroGNNv2(torch.nn.Module):
def __init__(self, hidden_channels, out_channels):
super().__init__()
self.conv1 = HeteroConv({
('user', 'rates', 'movie'): SAGEConv((-1, -1), hidden_channels),
('movie', 'rated_by', 'user'): SAGEConv((-1, -1), hidden_channels),
('movie', 'has_genre', 'genre'): SAGEConv((-1, -1), hidden_channels),
}, aggr='sum')
self.lin_user = Linear(hidden_channels, out_channels)
self.lin_movie = Linear(hidden_channels, out_channels)
def forward(self, x_dict, edge_index_dict):
x_dict = self.conv1(x_dict, edge_index_dict)
x_dict = {key: x.relu() for key, x in x_dict.items()}
return {
'user': self.lin_user(x_dict['user']),
'movie': self.lin_movie(x_dict['movie']),
}
model = HeteroGNNv2(hidden_channels=16, out_channels=4)
out_dict = model(data.x_dict, data.edge_index_dict)
print(f"User embeddings shape: {out_dict['user'].shape}")
print(f"Movie embeddings shape: {out_dict['movie'].shape}")
Discussion: The genre node type doesn't appear in out_dict because we never defined an output layer such as lin_genre for it, but the key point is that message passing through ('movie', 'has_genre', 'genre') lets genre information indirectly influence the movie node representations. In a heterogeneous graph, even node types that aren't used directly for prediction (here, genre) can be incorporated into the graph as an auxiliary source of information.
Summary
In this chapter, we learned about more advanced GNN architectures and graph-level tasks. Let's revisit the learning objectives from the beginning of the chapter.
- ✅ GAT: we understood how an attention mechanism lets the model learn a different importance for each node, achieving more flexible feature aggregation than GCN
- ✅ GraphSAGE: we learned the sampling strategy — sampling neighbors and training in mini-batches — that makes training feasible even on large graphs
- ✅ Graph pooling: we implemented global pooling functions such as
global_mean_pool, as well as hierarchical pooling withTopKPooling, to obtain a whole-graph representation - ✅ Graph classification: we built a pipeline that predicts molecular mutagenicity on the MUTAG dataset using a model combining GATConv and SAGEConv
- ✅ Link prediction: we understood the basics of predicting the probability that an edge exists, using an encoder-decoder architecture and negative sampling
- ✅ Heterogeneous graphs: we learned how to work with graphs containing multiple node and edge types using
HeteroDataandHeteroConv
🎉 Advanced GNN Architectures Mastered
Building on node classification (Chapter 2), you now have the ability to implement the major GNN tasks of graph classification, link prediction, and heterogeneous graphs. GAT, GraphSAGE, and pooling are among the most commonly combined techniques in real-world GNN applications. The next chapter applies these techniques to real-world data and digs deeper into efficiently handling large graphs, as well as model evaluation and tuning.
Reference Resources