How to return attention weights in multiple layers using Sequential and nn.ModuleList?

[fbragman]

Hi,

I am using torch_geometric.nn.Sequential with GATConv and I need to return the attention weights at multiple layers of the network. For dummy 1-layer network, I can do the following:

class Net(torch.nn.Module):
    def __init__(self, in_channels):
        super(Net, self).__init__()
        self.conv = Sequential(
            "x, edge_index, size, return_attention_weights",
            [
                (
                    GATConv(in_channels, 1, heads=2, dropout=0.6),
                    "x, edge_index, size, return_attention_weights -> x, a",
                ),
                (nn.ReLU(inplace=True), "x -> x"),
                (lambda x, a: [x, a], "x, a -> o")
            ],
        )

    def forward(self, x, edge_index):
        x = self.conv(x, edge_index, None, True)
        # returns output of GATConv/ReLU and attention weights all heads
        return x[0], x[1]

However, I am interested in building n-layer networks where n is a hyper-parameter and I would like to be able to output the attention weights at each layer. I have tried the following for a dummy 2 layer network

class Net(torch.nn.Module):
    def __init__(self, in_channels):
        super(Net, self).__init__()

        self.conv = torch.nn.ModuleList()

        self.conv.append(
            Sequential(
                "x, edge_index, size, return_attention_weights",
                [
                    (
                        GATConv(in_channels, 10, heads=2, dropout=0.6),
                        "x, edge_index, size, return_attention_weights -> x, a",
                    ),
                    (nn.ReLU(inplace=True), "x -> x"),
                    (lambda x, a: [x, a], "x, a -> o"),
                ],
            )
        )

        self.conv.append(
            Sequential(
                "x, edge_index, size, return_attention_weights",
                [
                    (
                        GATConv(20, 1, heads=1, dropout=0.6),
                        "x, edge_index, size, return_attention_weights -> x, a",
                    ),
                    (nn.ReLU(inplace=True), "x -> x"),
                    (lambda x, a: [x, a], "x, a -> o"),
                ],
            )
        )

    def forward(self, x, edge_index):
        atts = []
        for co in self.conv:
            out = co(x, edge_index, None, True)
            atts.append(out[1])
            x = out[0]
        return x

I think this works well but wondering if there is a better way?

Edit - I figured it out quickly after posting and thus edited the code posted in this discussion. However, open to suggestions on how to make this cleaner. For instance, seems I have to be specify Size in both Sequential and in the forward pass for each conv layer. Would be cleaner if I could just specify Sequential("x, edge_index, return_attention_weights") and conv(x, edge_index, True) without having to specify Size=None since it is defaulted to None in GATConv

[rusty1s]

One way to write this would be:

Sequential('x, edge_index, size, return_att', [
  (GATConv(...), 'x, edge_index, size, return_att' -> 'x, a1'),
  (ReLU(), 'x -> x'),
  (GATConv(...), 'x, edge_index, size, return_att' -> 'x, a2'),
  (ReLU(), 'x -> x'),
  (lambda x, a1, a2: [x, a1, a2], "x, a1, a2 -> o"),
])

However, I think such custom model is better defined without Sequential:

def forward(self, x, edge_index):
    x, a1 = self.conv1(x, edge_index, return_attention_weights=True)
    x= x.relu()
    x, a2 = self.conv1(x, edge_index, return_attention_weights=True)
    x= x.relu()
    return x, a1, a2

[songsong0425]

Hi, sorry for reusing this thread for the question. When I set the GAT model and tried to get the attention scores, my code returned the error. Where should I fix the code?

from torch_geometric.utils import remove_self_loops

# GAT construction
class GAT(torch.nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = GATConv(train_data.num_features, 64, heads=8) # in_channel, hidden_channel, heads
        self.conv2 = GATConv(64 * 8, 32, heads=1, concat=True) # hidden_channel * heads, out_channel, heads
        
    def forward(self, x, edge_index):
        x, a1 = self.conv1(x, edge_index, return_attention_weights=True)
        x = F.leaky_relu(x)
        x = F.dropout(x, p = 0.3, training = self.training)
        x, a2 = self.conv2(x, edge_index, return_attention_weights=True)

        return x, a1, a2
    
model = GAT().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-4)
criterion = torch.nn.BCEWithLogitsLoss()

...

def train(num_epochs):
    for epoch in range(1, num_epochs+1):
        tr_losses = 0
        val_losses = 0
        
        model.train()
        for data in tqdm(train_loader):
            data = data.to(device)
            optimizer.zero_grad()
        
            z = model(data.x, data.edge_index)
            
# Note that I run the code for the link prediction task
            out = ((z[data.edge_label_index[0]] * z[data.edge_label_index[1]]).sum(dim=-1)).view(-1)
            tr_loss = criterion(out, data.edge_label)
            tr_losses += tr_loss.item()
            tr_loss.backward()
            optimizer.step()
        avg_tr_loss = tr_losses/len(train_loader.dataset)
...

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[21], line 2
      1 num_epochs = 50
----> 2 train(num_epochs)

Cell In[19], line 20, in train(num_epochs)
     16 optimizer.zero_grad()
     18 z = model(data.x, data.edge_index)
---> 20 out = ((z[data.edge_label_index[0]] * z[data.edge_label_index[1]]).sum(dim=-1)).view(-1)
     21 tr_loss = criterion(out, data.edge_label)
     22 tr_losses += tr_loss.item()
TypeError: only integer tensors of a single element can be converted to an index

[songsong0425]

I simply solved the problem by changing the code z to z, a1, a2. Sorry for the unnecessary alarm!