Extract GAT Attention Weights when using NeighborSampler?

[jah377]

Hi,

I am interested in capturing attention weights computed during the GAT model. As a proof of concept, I extracted weights using the model outlined below. While successful on smaller datasets like Cora, this full batch approach is not feasible with larger datasets.

torch_geometric.loader.NeighborSampler is useful for dealing with larger datasets. PyG has a nice example employing GAT on larger datasets like OGB-Products. However, it is not obvious how attention weights would be extracted as you iterate over the loader.

Any help would be appreciated, Thanks!

class GAT(torch.nn.Module):
    def __init__(self, in_channels, out_channels, hidden_channels, nheads):

        super(GAT, self).__init__()

        self.conv1 = GATConv(in_channels, hidden_channels, nheads)
        self.conv2 = GATConv(hidden_channels*nheads, out_channels, nheads, concat=False)

    def extract_features(self, data, return_attention_weights=True):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x, edge_index)
        x = F.elu(x)
        x = F.dropout(x, p=0.6, training=self.training)
        return self.conv2(x, edge_index, return_attention_weights=return_attention_weights)

    def forward(self, data):
        x, _ = self.extract_features(data)
        return F.log_softmax(x, dim=1)

[rusty1s]

I think you can still use the same procedure as in full-batch mode for this (via out, (att_index, att_weight) = conv(x, edge_index)). The main challenge is to map the local att_index to global indices afterwards. This should be achievable via n_id[att_index].