Link prediction using GAT

[songsong0425]

Dear PyG community,
Before starting, thank you for your effort and sorry for frequent questions.

I'm trying to construct the GAT model for link prediction, but I countered an error.
This is the base code for the task:

from torch_geometric.utils import barabasi_albert_graph
from torch_geometric.loader import DataLoader
from torch_geometric.transforms import RandomLinkSplit
from torch_geometric.data import Data
from torch_geometric.nn import GATConv

# Example graph data generation
bagraph = barabasi_albert_graph(num_nodes=100, num_edges=80)
example_feature = torch.rand(100, 256)
example_graph = Data(edge_index=bagraph, x=example_feature)
example_graph

# train/val/test sepration
batch_size = 32
transform = RandomLinkSplit(num_val=0.1, num_test=0.2, is_undirected=True, split_labels=True)
train_data, val_data, test_data = transform(example_graph)

train_loader = DataLoader(train_data, shuffle=True, batch_size=batch_size)
val_loader = DataLoader(val_data, shuffle=True, batch_size=batch_size)
test_loader = DataLoader(test_data, shuffle=False, batch_size=batch_size)

# GAT construction
class GAT(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, heads):
        super(GAT, self).__init__()

        self.conv1 = GATConv(in_channels, hidden_channels, heads, dropout=0.3)
        self.conv2 = GATConv(hidden_channels * heads, out_channels, heads=1, concat=False, dropout=0.3)
        
    def forward(self, data):
        x = data.x
        
        edge_index = torch.cat([data.pos_edge_label_index, data.neg_edge_label_index], dim=-1)
        edge_label = torch.cat([data.pos_edge_label, data.neg_edge_label], dim=0)
        
        #x = F.dropout(x, p = 0.3, training = self.training)
        x = self.conv1(x, edge_index)
        x = F.elu(x)
        x = F.dropout(x, p = 0.3, training = self.training)
        x = self.conv2(x, edge_index)
        
        return x

model = GAT(example_graph.num_features, 64, 32, 8).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
criterion = torch.nn.BCEWithLogitsLoss()

# Define training process
def train():
    model.train()
    
    loss = 0
    for data in train_loader:
        optimizer.zero_grad()
    
        z = model(data)
        
        out = ((z[edge_index[0]] * z[edge_index[1]]).sum(dim=-1)).view(-1)
        loss = criterion(out, edge_label)
        loss.backward()
        optimizer.step()
    
    return float(loss)

# Test the model
@torch.no_grad()
def test(loader):
    model.eval()
    
    y_pred, y_true = [], []
    for data in loader:
        edge_index = torch.cat([data.pos_edge_label_index, data.neg_edge_label_index], dim=-1)
        edge_label = torch.cat([data.pos_edge_label, data.neg_edge_label], dim=0)
        
        z = model(data.x, data.edge_index)
        out = ((z[edge_index[0]] * z[edge_index[1]]).sum(dim=-1)).view(-1).sigmoid()
      
        return roc_auc_score(edge_label.cpu().numpy(), out.cpu().numpy)
    
best_val_auc = final_test_auc = 0

for epoch in range(1, 11):
    loss = train()
    val_auc = test(val_loader)
    test_auc = test(test_loader)
    
    if val_auc > best_val_auc:
        best_val_auc = val_auc
        final_test_auc = test_auc
    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Val: {val_auc:.4f}, '
          f'Test: {test_auc:.4f}')
    
print(f'Final Test: {final_test_auc:.4f}')

When I run the above code, the error said:
RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.
KeyError: 1

Full error message

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In [39], line 4
      1 best_val_auc = final_test_auc = 0
      3 for epoch in range(1, 51):
----> 4     loss = train()
      5     val_auc = test(val_loader)
      6     test_auc = test(test_loader)

Cell In [36], line 5, in train()
      2 model.train()
      4 loss = 0
----> 5 for data in train_loader:
      6     optimizer.zero_grad()
      8     z = model(data)

File ~/.venv/torch1.12.1-py3.8-cuda11.3/lib/python3.8/site-packages/torch/utils/data/dataloader.py:681, in _BaseDataLoaderIter.__next__(self)
    678 if self._sampler_iter is None:
    679     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    680     self._reset()  # type: ignore[call-arg]
--> 681 data = self._next_data()
    682 self._num_yielded += 1
    683 if self._dataset_kind == _DatasetKind.Iterable and \
    684         self._IterableDataset_len_called is not None and \
    685         self._num_yielded > self._IterableDataset_len_called:

File ~/.venv/torch1.12.1-py3.8-cuda11.3/lib/python3.8/site-packages/torch/utils/data/dataloader.py:721, in _SingleProcessDataLoaderIter._next_data(self)
    719 def _next_data(self):
    720     index = self._next_index()  # may raise StopIteration
--> 721     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    722     if self._pin_memory:
    723         data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)

File ~/.venv/torch1.12.1-py3.8-cuda11.3/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py:49, in _MapDatasetFetcher.fetch(self, possibly_batched_index)
     47 def fetch(self, possibly_batched_index):
     48     if self.auto_collation:
---> 49         data = [self.dataset[idx] for idx in possibly_batched_index]
     50     else:
     51         data = self.dataset[possibly_batched_index]

File ~/.venv/torch1.12.1-py3.8-cuda11.3/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py:49, in <listcomp>(.0)
     47 def fetch(self, possibly_batched_index):
     48     if self.auto_collation:
---> 49         data = [self.dataset[idx] for idx in possibly_batched_index]
     50     else:
     51         data = self.dataset[possibly_batched_index]

File ~/.venv/torch1.12.1-py3.8-cuda11.3/lib/python3.8/site-packages/torch_geometric/data/data.py:457, in Data.__getitem__(self, key)
    456 def __getitem__(self, key: str) -> Any:
--> 457     return self._store[key]

File ~/.venv/torch1.12.1-py3.8-cuda11.3/lib/python3.8/site-packages/torch_geometric/data/storage.py:104, in BaseStorage.__getitem__(self, key)
    103 def __getitem__(self, key: str) -> Any:
--> 104     return self._mapping[key]

KeyError: 1

References:

• GAT code (https://github.com/pyg-team/pytorch_geometric/blob/master/examples/gat.py)
• Link prediction code(https://github.com/pyg-team/pytorch_geometric/blob/master/examples/link_pred.py)

[rusty1s]

Does return float(loss) resolves this issue?

[songsong0425]

No, I'm debugging the code. Nevertheless, the reason I closed this discussion is that this question looks too naive and irresponsible. I tried your solution, but it still returned the same error. :(
I slightly modified the code and countered the new error.