Implementing a GAT in PyTorch Geometric

We now have a complete picture of how the graph attention layer works. These layers can be stacked to create our new architecture of choice: the GAT. In this section, we will follow the guidelines from the original GAT paper to implement our own model using PyG. We will use it to perform node classification on the Cora and CiteSeer datasets. Finally, we will comment on these results and compare them.

Let’s start with the Cora dataset:

1. We import Cora from the Planetoid class using PyG:

   from torch_geometric.datasets import Planetoid
   dataset = Planetoid(root=".", name="Cora")
   data = dataset[0]
   Data(x=[2708, 1433], edge_index=[2, 10556], y=[2708], train_mask=[2708], val_mask=[2708], test_mask=[2708])

2. We import the necessary libraries to create our own GAT class, using the GATv2 layer:

   import torch
   import torch.nn.functional as F
   from torch_geometric.nn import GATv2Conv
   from torch.nn import Linear, Dropout

...