wehub/labmlai--annotated_deep_learning_paper_implementations
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:19:01 +08:00
commit 3b90d1192f
2172 changed files with 594509 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
labml_nn/graphs/__init__.py
+12
View File
@@ -0,0 +1,12 @@
"""
---
title: Graph Neural Networks
summary: >
A set of PyTorch implementations/tutorials related to graph neural networks
---
# Graph Neural Networks
* [Graph Attention Networks (GAT)](gat/index.html)
* [Graph Attention Networks v2 (GATv2)](gatv2/index.html)
"""
labml_nn/graphs/gat/__init__.py
+196
View File
@@ -0,0 +1,196 @@
"""
---
title: Graph Attention Networks (GAT)
summary: >
A PyTorch implementation/tutorial of Graph Attention Networks.
---
# Graph Attention Networks (GAT)
This is a [PyTorch](https://pytorch.org) implementation of the paper
[Graph Attention Networks](https://arxiv.org/abs/1710.10903).
GATs work on graph data.
A graph consists of nodes and edges connecting nodes.
For example, in Cora dataset the nodes are research papers and the edges are citations that
connect the papers.
GAT uses masked self-attention, kind of similar to [transformers](../../transformers/mha.html).
GAT consists of graph attention layers stacked on top of each other.
Each graph attention layer gets node embeddings as inputs and outputs transformed embeddings.
The node embeddings pay attention to the embeddings of other nodes it's connected to.
The details of graph attention layers are included alongside the implementation.
Here is [the training code](experiment.html) for training
a two-layer GAT on Cora dataset.
"""
import torch
from torch import nn
class GraphAttentionLayer(nn.Module):
"""
## Graph attention layer
This is a single graph attention layer.
A GAT is made up of multiple such layers.
It takes
$$\mathbf{h} = \{ \overrightarrow{h_1}, \overrightarrow{h_2}, \dots, \overrightarrow{h_N} \}$$,
where $\overrightarrow{h_i} \in \mathbb{R}^F$ as input
and outputs
$$\mathbf{h'} = \{ \overrightarrow{h'_1}, \overrightarrow{h'_2}, \dots, \overrightarrow{h'_N} \}$$,
where $\overrightarrow{h'_i} \in \mathbb{R}^{F'}$.
"""
def __init__(self, in_features: int, out_features: int, n_heads: int,
is_concat: bool = True,
dropout: float = 0.6,
leaky_relu_negative_slope: float = 0.2):
"""
* `in_features`, $F$, is the number of input features per node
* `out_features`, $F'$, is the number of output features per node
* `n_heads`, $K$, is the number of attention heads
* `is_concat` whether the multi-head results should be concatenated or averaged
* `dropout` is the dropout probability
* `leaky_relu_negative_slope` is the negative slope for leaky relu activation
"""
super().__init__()
self.is_concat = is_concat
self.n_heads = n_heads
# Calculate the number of dimensions per head
if is_concat:
assert out_features % n_heads == 0
# If we are concatenating the multiple heads
self.n_hidden = out_features // n_heads
else:
# If we are averaging the multiple heads
self.n_hidden = out_features
# Linear layer for initial transformation;
# i.e. to transform the node embeddings before self-attention
self.linear = nn.Linear(in_features, self.n_hidden * n_heads, bias=False)
# Linear layer to compute attention score $e_{ij}$
self.attn = nn.Linear(self.n_hidden * 2, 1, bias=False)
# The activation for attention score $e_{ij}$
self.activation = nn.LeakyReLU(negative_slope=leaky_relu_negative_slope)
# Softmax to compute attention $\alpha_{ij}$
self.softmax = nn.Softmax(dim=1)
# Dropout layer to be applied for attention
self.dropout = nn.Dropout(dropout)
def forward(self, h: torch.Tensor, adj_mat: torch.Tensor):
"""
* `h`, $\mathbf{h}$ is the input node embeddings of shape `[n_nodes, in_features]`.
* `adj_mat` is the adjacency matrix of shape `[n_nodes, n_nodes, n_heads]`.
We use shape `[n_nodes, n_nodes, 1]` since the adjacency is the same for each head.
Adjacency matrix represent the edges (or connections) among nodes.
`adj_mat[i][j]` is `True` if there is an edge from node `i` to node `j`.
"""
# Number of nodes
n_nodes = h.shape[0]
# The initial transformation,
# $$\overrightarrow{g^k_i} = \mathbf{W}^k \overrightarrow{h_i}$$
# for each head.
# We do single linear transformation and then split it up for each head.
g = self.linear(h).view(n_nodes, self.n_heads, self.n_hidden)
# #### Calculate attention score
#
# We calculate these for each head $k$. *We have omitted $\cdot^k$ for simplicity*.
#
# $$e_{ij} = a(\mathbf{W} \overrightarrow{h_i}, \mathbf{W} \overrightarrow{h_j}) =
# a(\overrightarrow{g_i}, \overrightarrow{g_j})$$
#
# $e_{ij}$ is the attention score (importance) from node $j$ to node $i$.
# We calculate this for each head.
#
# $a$ is the attention mechanism, that calculates the attention score.
# The paper concatenates
# $\overrightarrow{g_i}$, $\overrightarrow{g_j}$
# and does a linear transformation with a weight vector $\mathbf{a} \in \mathbb{R}^{2 F'}$
# followed by a $\text{LeakyReLU}$.
#
# $$e_{ij} = \text{LeakyReLU} \Big(
# \mathbf{a}^\top \Big[
# \overrightarrow{g_i} \Vert \overrightarrow{g_j}
# \Big] \Big)$$
# First we calculate
# $\Big[\overrightarrow{g_i} \Vert \overrightarrow{g_j} \Big]$
# for all pairs of $i, j$.
#
# `g_repeat` gets
# $$\{\overrightarrow{g_1}, \overrightarrow{g_2}, \dots, \overrightarrow{g_N},
# \overrightarrow{g_1}, \overrightarrow{g_2}, \dots, \overrightarrow{g_N}, ...\}$$
# where each node embedding is repeated `n_nodes` times.
g_repeat = g.repeat(n_nodes, 1, 1)
# `g_repeat_interleave` gets
# $$\{\overrightarrow{g_1}, \overrightarrow{g_1}, \dots, \overrightarrow{g_1},
# \overrightarrow{g_2}, \overrightarrow{g_2}, \dots, \overrightarrow{g_2}, ...\}$$
# where each node embedding is repeated `n_nodes` times.
g_repeat_interleave = g.repeat_interleave(n_nodes, dim=0)
# Now we concatenate to get
# $$\{\overrightarrow{g_1} \Vert \overrightarrow{g_1},
# \overrightarrow{g_1} \Vert \overrightarrow{g_2},
# \dots, \overrightarrow{g_1} \Vert \overrightarrow{g_N},
# \overrightarrow{g_2} \Vert \overrightarrow{g_1},
# \overrightarrow{g_2} \Vert \overrightarrow{g_2},
# \dots, \overrightarrow{g_2} \Vert \overrightarrow{g_N}, ...\}$$
g_concat = torch.cat([g_repeat_interleave, g_repeat], dim=-1)
# Reshape so that `g_concat[i, j]` is $\overrightarrow{g_i} \Vert \overrightarrow{g_j}$
g_concat = g_concat.view(n_nodes, n_nodes, self.n_heads, 2 * self.n_hidden)
# Calculate
# $$e_{ij} = \text{LeakyReLU} \Big(
# \mathbf{a}^\top \Big[
# \overrightarrow{g_i} \Vert \overrightarrow{g_j}
# \Big] \Big)$$
# `e` is of shape `[n_nodes, n_nodes, n_heads, 1]`
e = self.activation(self.attn(g_concat))
# Remove the last dimension of size `1`
e = e.squeeze(-1)
# The adjacency matrix should have shape
# `[n_nodes, n_nodes, n_heads]` or`[n_nodes, n_nodes, 1]`
assert adj_mat.shape[0] == 1 or adj_mat.shape[0] == n_nodes
assert adj_mat.shape[1] == 1 or adj_mat.shape[1] == n_nodes
assert adj_mat.shape[2] == 1 or adj_mat.shape[2] == self.n_heads
# Mask $e_{ij}$ based on adjacency matrix.
# $e_{ij}$ is set to $- \infty$ if there is no edge from $i$ to $j$.
e = e.masked_fill(adj_mat == 0, float('-inf'))
# We then normalize attention scores (or coefficients)
# $$\alpha_{ij} = \text{softmax}_j(e_{ij}) =
# \frac{\exp(e_{ij})}{\sum_{k \in \mathcal{N}_i} \exp(e_{ik})}$$
#
# where $\mathcal{N}_i$ is the set of nodes connected to $i$.
#
# We do this by setting unconnected $e_{ij}$ to $- \infty$ which
# makes $\exp(e_{ij}) \sim 0$ for unconnected pairs.
a = self.softmax(e)
# Apply dropout regularization
a = self.dropout(a)
# Calculate final output for each head
# $$\overrightarrow{h'^k_i} = \sum_{j \in \mathcal{N}_i} \alpha^k_{ij} \overrightarrow{g^k_j}$$
#
# *Note:* The paper includes the final activation $\sigma$ in $\overrightarrow{h_i}$
# We have omitted this from the Graph Attention Layer implementation
# and use it on the GAT model to match with how other PyTorch modules are defined -
# activation as a separate layer.
attn_res = torch.einsum('ijh,jhf->ihf', a, g)
# Concatenate the heads
if self.is_concat:
# $$\overrightarrow{h'_i} = \Bigg\Vert_{k=1}^{K} \overrightarrow{h'^k_i}$$
return attn_res.reshape(n_nodes, self.n_heads * self.n_hidden)
# Take the mean of the heads
else:
# $$\overrightarrow{h'_i} = \frac{1}{K} \sum_{k=1}^{K} \overrightarrow{h'^k_i}$$
return attn_res.mean(dim=1)
labml_nn/graphs/gat/experiment.py
+309
View File
@@ -0,0 +1,309 @@
"""
---
title: Train a Graph Attention Network (GAT) on Cora dataset
summary: >
This trains is a Graph Attention Network (GAT) on Cora dataset
---
# Train a Graph Attention Network (GAT) on Cora dataset
"""
from typing import Dict
import numpy as np
import torch
from torch import nn
from labml import lab, monit, tracker, experiment
from labml.configs import BaseConfigs, option, calculate
from labml.utils import download
from labml_nn.helpers.device import DeviceConfigs
from labml_nn.graphs.gat import GraphAttentionLayer
from labml_nn.optimizers.configs import OptimizerConfigs
class CoraDataset:
"""
## [Cora Dataset](https://linqs.soe.ucsc.edu/data)
Cora dataset is a dataset of research papers.
For each paper we are given a binary feature vector that indicates the presence of words.
Each paper is classified into one of 7 classes.
The dataset also has the citation network.
The papers are the nodes of the graph and the edges are the citations.
The task is to classify the nodes to the 7 classes with feature vectors and
citation network as input.
"""
# Labels for each node
labels: torch.Tensor
# Set of class names and an unique integer index
classes: Dict[str, int]
# Feature vectors for all nodes
features: torch.Tensor
# Adjacency matrix with the edge information.
# `adj_mat[i][j]` is `True` if there is an edge from `i` to `j`.
adj_mat: torch.Tensor
@staticmethod
def _download():
"""
Download the dataset
"""
if not (lab.get_data_path() / 'cora').exists():
download.download_file('https://linqs-data.soe.ucsc.edu/public/lbc/cora.tgz',
lab.get_data_path() / 'cora.tgz')
download.extract_tar(lab.get_data_path() / 'cora.tgz', lab.get_data_path())
def __init__(self, include_edges: bool = True):
"""
Load the dataset
"""
# Whether to include edges.
# This is test how much accuracy is lost if we ignore the citation network.
self.include_edges = include_edges
# Download dataset
self._download()
# Read the paper ids, feature vectors, and labels
with monit.section('Read content file'):
content = np.genfromtxt(str(lab.get_data_path() / 'cora/cora.content'), dtype=np.dtype(str))
# Load the citations, it's a list of pairs of integers.
with monit.section('Read citations file'):
citations = np.genfromtxt(str(lab.get_data_path() / 'cora/cora.cites'), dtype=np.int32)
# Get the feature vectors
features = torch.tensor(np.array(content[:, 1:-1], dtype=np.float32))
# Normalize the feature vectors
self.features = features / features.sum(dim=1, keepdim=True)
# Get the class names and assign an unique integer to each of them
self.classes = {s: i for i, s in enumerate(set(content[:, -1]))}
# Get the labels as those integers
self.labels = torch.tensor([self.classes[i] for i in content[:, -1]], dtype=torch.long)
# Get the paper ids
paper_ids = np.array(content[:, 0], dtype=np.int32)
# Map of paper id to index
ids_to_idx = {id_: i for i, id_ in enumerate(paper_ids)}
# Empty adjacency matrix - an identity matrix
self.adj_mat = torch.eye(len(self.labels), dtype=torch.bool)
# Mark the citations in the adjacency matrix
if self.include_edges:
for e in citations:
# The pair of paper indexes
e1, e2 = ids_to_idx[e[0]], ids_to_idx[e[1]]
# We build a symmetrical graph, where if paper $i$ referenced
# paper $j$ we place an adge from $i$ to $j$ as well as an edge
# from $j$ to $i$.
self.adj_mat[e1][e2] = True
self.adj_mat[e2][e1] = True
class GAT(nn.Module):
"""
## Graph Attention Network (GAT)
This graph attention network has two [graph attention layers](index.html).
"""
def __init__(self, in_features: int, n_hidden: int, n_classes: int, n_heads: int, dropout: float):
"""
* `in_features` is the number of features per node
* `n_hidden` is the number of features in the first graph attention layer
* `n_classes` is the number of classes
* `n_heads` is the number of heads in the graph attention layers
* `dropout` is the dropout probability
"""
super().__init__()
# First graph attention layer where we concatenate the heads
self.layer1 = GraphAttentionLayer(in_features, n_hidden, n_heads, is_concat=True, dropout=dropout)
# Activation function after first graph attention layer
self.activation = nn.ELU()
# Final graph attention layer where we average the heads
self.output = GraphAttentionLayer(n_hidden, n_classes, 1, is_concat=False, dropout=dropout)
# Dropout
self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor, adj_mat: torch.Tensor):
"""
* `x` is the features vectors of shape `[n_nodes, in_features]`
* `adj_mat` is the adjacency matrix of the form
`[n_nodes, n_nodes, n_heads]` or `[n_nodes, n_nodes, 1]`
"""
# Apply dropout to the input
x = self.dropout(x)
# First graph attention layer
x = self.layer1(x, adj_mat)
# Activation function
x = self.activation(x)
# Dropout
x = self.dropout(x)
# Output layer (without activation) for logits
return self.output(x, adj_mat)
def accuracy(output: torch.Tensor, labels: torch.Tensor):
"""
A simple function to calculate the accuracy
"""
return output.argmax(dim=-1).eq(labels).sum().item() / len(labels)
class Configs(BaseConfigs):
"""
## Configurations
"""
# Model
model: GAT
# Number of nodes to train on
training_samples: int = 500
# Number of features per node in the input
in_features: int
# Number of features in the first graph attention layer
n_hidden: int = 64
# Number of heads
n_heads: int = 8
# Number of classes for classification
n_classes: int
# Dropout probability
dropout: float = 0.6
# Whether to include the citation network
include_edges: bool = True
# Dataset
dataset: CoraDataset
# Number of training iterations
epochs: int = 1_000
# Loss function
loss_func = nn.CrossEntropyLoss()
# Device to train on
#
# This creates configs for device, so that
# we can change the device by passing a config value
device: torch.device = DeviceConfigs()
# Optimizer
optimizer: torch.optim.Adam
def run(self):
"""
### Training loop
We do full batch training since the dataset is small.
If we were to sample and train we will have to sample a set of
nodes for each training step along with the edges that span
across those selected nodes.
"""
# Move the feature vectors to the device
features = self.dataset.features.to(self.device)
# Move the labels to the device
labels = self.dataset.labels.to(self.device)
# Move the adjacency matrix to the device
edges_adj = self.dataset.adj_mat.to(self.device)
# Add an empty third dimension for the heads
edges_adj = edges_adj.unsqueeze(-1)
# Random indexes
idx_rand = torch.randperm(len(labels))
# Nodes for training
idx_train = idx_rand[:self.training_samples]
# Nodes for validation
idx_valid = idx_rand[self.training_samples:]
# Training loop
for epoch in monit.loop(self.epochs):
# Set the model to training mode
self.model.train()
# Make all the gradients zero
self.optimizer.zero_grad()
# Evaluate the model
output = self.model(features, edges_adj)
# Get the loss for training nodes
loss = self.loss_func(output[idx_train], labels[idx_train])
# Calculate gradients
loss.backward()
# Take optimization step
self.optimizer.step()
# Log the loss
tracker.add('loss.train', loss)
# Log the accuracy
tracker.add('accuracy.train', accuracy(output[idx_train], labels[idx_train]))
# Set mode to evaluation mode for validation
self.model.eval()
# No need to compute gradients
with torch.no_grad():
# Evaluate the model again
output = self.model(features, edges_adj)
# Calculate the loss for validation nodes
loss = self.loss_func(output[idx_valid], labels[idx_valid])
# Log the loss
tracker.add('loss.valid', loss)
# Log the accuracy
tracker.add('accuracy.valid', accuracy(output[idx_valid], labels[idx_valid]))
# Save logs
tracker.save()
@option(Configs.dataset)
def cora_dataset(c: Configs):
"""
Create Cora dataset
"""
return CoraDataset(c.include_edges)
# Get the number of classes
calculate(Configs.n_classes, lambda c: len(c.dataset.classes))
# Number of features in the input
calculate(Configs.in_features, lambda c: c.dataset.features.shape[1])
@option(Configs.model)
def gat_model(c: Configs):
"""
Create GAT model
"""
return GAT(c.in_features, c.n_hidden, c.n_classes, c.n_heads, c.dropout).to(c.device)
@option(Configs.optimizer)
def _optimizer(c: Configs):
"""
Create configurable optimizer
"""
opt_conf = OptimizerConfigs()
opt_conf.parameters = c.model.parameters()
return opt_conf
def main():
# Create configurations
conf = Configs()
# Create an experiment
experiment.create(name='gat')
# Calculate configurations.
experiment.configs(conf, {
# Adam optimizer
'optimizer.optimizer': 'Adam',
'optimizer.learning_rate': 5e-3,
'optimizer.weight_decay': 5e-4,
})
# Start and watch the experiment
with experiment.start():
# Run the training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/graphs/gat/readme.md
+18
View File
@@ -0,0 +1,18 @@
# [Graph Attention Networks (GAT)](https://nn.labml.ai/graphs/gat/index.html)
This is a [PyTorch](https://pytorch.org) implementation of the paper
[Graph Attention Networks](https://arxiv.org/abs/1710.10903).
GATs work on graph data.
A graph consists of nodes and edges connecting nodes.
For example, in Cora dataset the nodes are research papers and the edges are citations that
connect the papers.
GAT uses masked self-attention, kind of similar to [transformers](https://nn.labml.ai/transformers/mha.html).
GAT consists of graph attention layers stacked on top of each other.
Each graph attention layer gets node embeddings as inputs and outputs transformed embeddings.
The node embeddings pay attention to the embeddings of other nodes it's connected to.
The details of graph attention layers are included alongside the implementation.
Here is [the training code](https://nn.labml.ai/graphs/gat/experiment.html) for training
a two-layer GAT on Cora dataset.
labml_nn/graphs/gatv2/__init__.py
+236
View File
@@ -0,0 +1,236 @@
"""
---
title: Graph Attention Networks v2 (GATv2)
summary: >
A PyTorch implementation/tutorial of Graph Attention Networks v2.
---
# Graph Attention Networks v2 (GATv2)
This is a [PyTorch](https://pytorch.org) implementation of the GATv2 operator from the paper
[How Attentive are Graph Attention Networks?](https://arxiv.org/abs/2105.14491).
GATv2s work on graph data similar to [GAT](../gat/index.html).
A graph consists of nodes and edges connecting nodes.
For example, in Cora dataset the nodes are research papers and the edges are citations that
connect the papers.
The GATv2 operator fixes the static attention problem of the standard [GAT](../gat/index.html).
Static attention is when the attention to the key nodes has the same rank (order) for any query node.
[GAT](../gat/index.html) computes attention from query node $i$ to key node $j$ as,
\begin{align}
e_{ij} &= \text{LeakyReLU} \Big(\mathbf{a}^\top \Big[
\mathbf{W} \overrightarrow{h_i} \Vert \mathbf{W} \overrightarrow{h_j}
\Big] \Big) \\
&=
\text{LeakyReLU} \Big(\mathbf{a}_1^\top \mathbf{W} \overrightarrow{h_i} +
\mathbf{a}_2^\top \mathbf{W} \overrightarrow{h_j}
\Big)
\end{align}
Note that for any query node $i$, the attention rank ($argsort$) of keys depends only
on $\mathbf{a}_2^\top \mathbf{W} \overrightarrow{h_j}$.
Therefore the attention rank of keys remains the same (*static*) for all queries.
GATv2 allows dynamic attention by changing the attention mechanism,
\begin{align}
e_{ij} &= \mathbf{a}^\top \text{LeakyReLU} \Big( \mathbf{W} \Big[
\overrightarrow{h_i} \Vert \overrightarrow{h_j}
\Big] \Big) \\
&= \mathbf{a}^\top \text{LeakyReLU} \Big(
\mathbf{W}_l \overrightarrow{h_i} + \mathbf{W}_r \overrightarrow{h_j}
\Big)
\end{align}
The paper shows that GATs static attention mechanism fails on some graph problems
with a synthetic dictionary lookup dataset.
It's a fully connected bipartite graph where one set of nodes (query nodes)
have a key associated with it
and the other set of nodes have both a key and a value associated with it.
The goal is to predict the values of query nodes.
GAT fails on this task because of its limited static attention.
Here is [the training code](experiment.html) for training
a two-layer GATv2 on Cora dataset.
"""
import torch
from torch import nn
class GraphAttentionV2Layer(nn.Module):
"""
## Graph attention v2 layer
This is a single graph attention v2 layer.
A GATv2 is made up of multiple such layers.
It takes
$$\mathbf{h} = \{ \overrightarrow{h_1}, \overrightarrow{h_2}, \dots, \overrightarrow{h_N} \}$$,
where $\overrightarrow{h_i} \in \mathbb{R}^F$ as input
and outputs
$$\mathbf{h'} = \{ \overrightarrow{h'_1}, \overrightarrow{h'_2}, \dots, \overrightarrow{h'_N} \}$$,
where $\overrightarrow{h'_i} \in \mathbb{R}^{F'}$.
"""
def __init__(self, in_features: int, out_features: int, n_heads: int,
is_concat: bool = True,
dropout: float = 0.6,
leaky_relu_negative_slope: float = 0.2,
share_weights: bool = False):
"""
* `in_features`, $F$, is the number of input features per node
* `out_features`, $F'$, is the number of output features per node
* `n_heads`, $K$, is the number of attention heads
* `is_concat` whether the multi-head results should be concatenated or averaged
* `dropout` is the dropout probability
* `leaky_relu_negative_slope` is the negative slope for leaky relu activation
* `share_weights` if set to `True`, the same matrix will be applied to the source and the target node of every edge
"""
super().__init__()
self.is_concat = is_concat
self.n_heads = n_heads
self.share_weights = share_weights
# Calculate the number of dimensions per head
if is_concat:
assert out_features % n_heads == 0
# If we are concatenating the multiple heads
self.n_hidden = out_features // n_heads
else:
# If we are averaging the multiple heads
self.n_hidden = out_features
# Linear layer for initial source transformation;
# i.e. to transform the source node embeddings before self-attention
self.linear_l = nn.Linear(in_features, self.n_hidden * n_heads, bias=False)
# If `share_weights` is `True` the same linear layer is used for the target nodes
if share_weights:
self.linear_r = self.linear_l
else:
self.linear_r = nn.Linear(in_features, self.n_hidden * n_heads, bias=False)
# Linear layer to compute attention score $e_{ij}$
self.attn = nn.Linear(self.n_hidden, 1, bias=False)
# The activation for attention score $e_{ij}$
self.activation = nn.LeakyReLU(negative_slope=leaky_relu_negative_slope)
# Softmax to compute attention $\alpha_{ij}$
self.softmax = nn.Softmax(dim=1)
# Dropout layer to be applied for attention
self.dropout = nn.Dropout(dropout)
def forward(self, h: torch.Tensor, adj_mat: torch.Tensor):
"""
* `h`, $\mathbf{h}$ is the input node embeddings of shape `[n_nodes, in_features]`.
* `adj_mat` is the adjacency matrix of shape `[n_nodes, n_nodes, n_heads]`.
We use shape `[n_nodes, n_nodes, 1]` since the adjacency is the same for each head.
Adjacency matrix represent the edges (or connections) among nodes.
`adj_mat[i][j]` is `True` if there is an edge from node `i` to node `j`.
"""
# Number of nodes
n_nodes = h.shape[0]
# The initial transformations,
# $$\overrightarrow{{g_l}^k_i} = \mathbf{W_l}^k \overrightarrow{h_i}$$
# $$\overrightarrow{{g_r}^k_i} = \mathbf{W_r}^k \overrightarrow{h_i}$$
# for each head.
# We do two linear transformations and then split it up for each head.
g_l = self.linear_l(h).view(n_nodes, self.n_heads, self.n_hidden)
g_r = self.linear_r(h).view(n_nodes, self.n_heads, self.n_hidden)
# #### Calculate attention score
#
# We calculate these for each head $k$. *We have omitted $\cdot^k$ for simplicity*.
#
# $$e_{ij} = a(\mathbf{W_l} \overrightarrow{h_i}, \mathbf{W_r} \overrightarrow{h_j}) =
# a(\overrightarrow{{g_l}_i}, \overrightarrow{{g_r}_j})$$
#
# $e_{ij}$ is the attention score (importance) from node $j$ to node $i$.
# We calculate this for each head.
#
# $a$ is the attention mechanism, that calculates the attention score.
# The paper sums
# $\overrightarrow{{g_l}_i}$, $\overrightarrow{{g_r}_j}$
# followed by a $\text{LeakyReLU}$
# and does a linear transformation with a weight vector $\mathbf{a} \in \mathbb{R}^{F'}$
#
#
# $$e_{ij} = \mathbf{a}^\top \text{LeakyReLU} \Big(
# \Big[
# \overrightarrow{{g_l}_i} + \overrightarrow{{g_r}_j}
# \Big] \Big)$$
# Note: The paper desrcibes $e_{ij}$ as
# $$e_{ij} = \mathbf{a}^\top \text{LeakyReLU} \Big( \mathbf{W}
# \Big[
# \overrightarrow{h_i} \Vert \overrightarrow{h_j}
# \Big] \Big)$$
# which is equivalent to the definition we use here.
# First we calculate
# $\Big[\overrightarrow{{g_l}_i} + \overrightarrow{{g_r}_j} \Big]$
# for all pairs of $i, j$.
#
# `g_l_repeat` gets
# $$\{\overrightarrow{{g_l}_1}, \overrightarrow{{g_l}_2}, \dots, \overrightarrow{{g_l}_N},
# \overrightarrow{{g_l}_1}, \overrightarrow{{g_l}_2}, \dots, \overrightarrow{{g_l}_N}, ...\}$$
# where each node embedding is repeated `n_nodes` times.
g_l_repeat = g_l.repeat(n_nodes, 1, 1)
# `g_r_repeat_interleave` gets
# $$\{\overrightarrow{{g_r}_1}, \overrightarrow{{g_r}_1}, \dots, \overrightarrow{{g_r}_1},
# \overrightarrow{{g_r}_2}, \overrightarrow{{g_r}_2}, \dots, \overrightarrow{{g_r}_2}, ...\}$$
# where each node embedding is repeated `n_nodes` times.
g_r_repeat_interleave = g_r.repeat_interleave(n_nodes, dim=0)
# Now we add the two tensors to get
# $$\{\overrightarrow{{g_l}_1} + \overrightarrow{{g_r}_1},
# \overrightarrow{{g_l}_1} + \overrightarrow{{g_r}_2},
# \dots, \overrightarrow{{g_l}_1} +\overrightarrow{{g_r}_N},
# \overrightarrow{{g_l}_2} + \overrightarrow{{g_r}_1},
# \overrightarrow{{g_l}_2} + \overrightarrow{{g_r}_2},
# \dots, \overrightarrow{{g_l}_2} + \overrightarrow{{g_r}_N}, ...\}$$
g_sum = g_l_repeat + g_r_repeat_interleave
# Reshape so that `g_sum[i, j]` is $\overrightarrow{{g_l}_i} + \overrightarrow{{g_r}_j}$
g_sum = g_sum.view(n_nodes, n_nodes, self.n_heads, self.n_hidden)
# Calculate
# $$e_{ij} = \mathbf{a}^\top \text{LeakyReLU} \Big(
# \Big[
# \overrightarrow{{g_l}_i} + \overrightarrow{{g_r}_j}
# \Big] \Big)$$
# `e` is of shape `[n_nodes, n_nodes, n_heads, 1]`
e = self.attn(self.activation(g_sum))
# Remove the last dimension of size `1`
e = e.squeeze(-1)
# The adjacency matrix should have shape
# `[n_nodes, n_nodes, n_heads]` or`[n_nodes, n_nodes, 1]`
assert adj_mat.shape[0] == 1 or adj_mat.shape[0] == n_nodes
assert adj_mat.shape[1] == 1 or adj_mat.shape[1] == n_nodes
assert adj_mat.shape[2] == 1 or adj_mat.shape[2] == self.n_heads
# Mask $e_{ij}$ based on adjacency matrix.
# $e_{ij}$ is set to $- \infty$ if there is no edge from $i$ to $j$.
e = e.masked_fill(adj_mat == 0, float('-inf'))
# We then normalize attention scores (or coefficients)
# $$\alpha_{ij} = \text{softmax}_j(e_{ij}) =
# \frac{\exp(e_{ij})}{\sum_{j' \in \mathcal{N}_i} \exp(e_{ij'})}$$
#
# where $\mathcal{N}_i$ is the set of nodes connected to $i$.
#
# We do this by setting unconnected $e_{ij}$ to $- \infty$ which
# makes $\exp(e_{ij}) \sim 0$ for unconnected pairs.
a = self.softmax(e)
# Apply dropout regularization
a = self.dropout(a)
# Calculate final output for each head
# $$\overrightarrow{h'^k_i} = \sum_{j \in \mathcal{N}_i} \alpha^k_{ij} \overrightarrow{{g_r}_{j,k}}$$
attn_res = torch.einsum('ijh,jhf->ihf', a, g_r)
# Concatenate the heads
if self.is_concat:
# $$\overrightarrow{h'_i} = \Bigg\Vert_{k=1}^{K} \overrightarrow{h'^k_i}$$
return attn_res.reshape(n_nodes, self.n_heads * self.n_hidden)
# Take the mean of the heads
else:
# $$\overrightarrow{h'_i} = \frac{1}{K} \sum_{k=1}^{K} \overrightarrow{h'^k_i}$$
return attn_res.mean(dim=1)
labml_nn/graphs/gatv2/experiment.py
+113
View File
@@ -0,0 +1,113 @@
"""
---
title: Train a Graph Attention Network v2 (GATv2) on Cora dataset
summary: >
This trains is a Graph Attention Network v2 (GATv2) on Cora dataset
---
# Train a Graph Attention Network v2 (GATv2) on Cora dataset
"""
import torch
from torch import nn
from labml import experiment
from labml.configs import option
from labml_nn.graphs.gat.experiment import Configs as GATConfigs
from labml_nn.graphs.gatv2 import GraphAttentionV2Layer
class GATv2(nn.Module):
"""
## Graph Attention Network v2 (GATv2)
This graph attention network has two [graph attention layers](index.html).
"""
def __init__(self, in_features: int, n_hidden: int, n_classes: int, n_heads: int, dropout: float,
share_weights: bool = True):
"""
* `in_features` is the number of features per node
* `n_hidden` is the number of features in the first graph attention layer
* `n_classes` is the number of classes
* `n_heads` is the number of heads in the graph attention layers
* `dropout` is the dropout probability
* `share_weights` if set to True, the same matrix will be applied to the source and the target node of every edge
"""
super().__init__()
# First graph attention layer where we concatenate the heads
self.layer1 = GraphAttentionV2Layer(in_features, n_hidden, n_heads,
is_concat=True, dropout=dropout, share_weights=share_weights)
# Activation function after first graph attention layer
self.activation = nn.ELU()
# Final graph attention layer where we average the heads
self.output = GraphAttentionV2Layer(n_hidden, n_classes, 1,
is_concat=False, dropout=dropout, share_weights=share_weights)
# Dropout
self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor, adj_mat: torch.Tensor):
"""
* `x` is the features vectors of shape `[n_nodes, in_features]`
* `adj_mat` is the adjacency matrix of the form
`[n_nodes, n_nodes, n_heads]` or `[n_nodes, n_nodes, 1]`
"""
# Apply dropout to the input
x = self.dropout(x)
# First graph attention layer
x = self.layer1(x, adj_mat)
# Activation function
x = self.activation(x)
# Dropout
x = self.dropout(x)
# Output layer (without activation) for logits
return self.output(x, adj_mat)
class Configs(GATConfigs):
"""
## Configurations
Since the experiment is same as [GAT experiment](../gat/experiment.html) but with
[GATv2 model](index.html) we extend the same configs and change the model.
"""
# Whether to share weights for source and target nodes of edges
share_weights: bool = False
# Set the model
model: GATv2 = 'gat_v2_model'
@option(Configs.model)
def gat_v2_model(c: Configs):
"""
Create GATv2 model
"""
return GATv2(c.in_features, c.n_hidden, c.n_classes, c.n_heads, c.dropout, c.share_weights).to(c.device)
def main():
# Create configurations
conf = Configs()
# Create an experiment
experiment.create(name='gatv2')
# Calculate configurations.
experiment.configs(conf, {
# Adam optimizer
'optimizer.optimizer': 'Adam',
'optimizer.learning_rate': 5e-3,
'optimizer.weight_decay': 5e-4,
'dropout': 0.7,
})
# Start and watch the experiment
with experiment.start():
# Run the training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/graphs/gatv2/readme.md
+17
View File
@@ -0,0 +1,17 @@
# [Graph Attention Networks v2 (GATv2)](https://nn.labml.ai/graphs/gatv2/index.html)
This is a [PyTorch](https://pytorch.org) implementation of the GATv2 operator from the paper
[How Attentive are Graph Attention Networks?](https://arxiv.org/abs/2105.14491).
GATv2s work on graph data.
A graph consists of nodes and edges connecting nodes.
For example, in Cora dataset the nodes are research papers and the edges are citations that
connect the papers.
The GATv2 operator fixes the static attention problem of the standard GAT:
since the linear layers in the standard GAT are applied right after each other, the ranking
of attended nodes is unconditioned on the query node.
In contrast, in GATv2, every node can attend to any other node.
Here is [the training code](https://nn.labml.ai/graphs/gatv2/experiment.html) for training
a two-layer GATv2 on Cora dataset.