How to create attention maps from attention weights from GATConv?

I installed dgl from source, and the GATConv.forward() has a get_attention parameter. This seems to return the α (attention coefficient) values of each edge from the neighborhood of a node. These coefficients are returned as a tensor. How can I create an attention map that can highlight the most important nodes / the relation between nodes using these weights?

You can do something as follows:

import dgl
import torch
from scipy.sparse import coo_matrix

# Assume eweight is the attention tensor returned from a `GATConv` instance.
eweight = torch.randn(3, 2, 1) # 3 edges, 2 heads
g = dgl.graph(([0, 1, 2], [1, 0, 1]))
num_nodes = g.num_nodes()
src, dst = g.edges(order='eid', form='uv')
edges = torch.stack([dst, src], dim=0)
attention_adjs = []
num_heads = eweight.shape[1]
for head in range(num_heads):
    atten_head_adj = coo_matrix((eweight[:, head, 0], (dst, src)), shape=(num_nodes, num_nodes))
    attention_adjs.append(atten_head_adj)

Once you have processed attention_adjs, you can then perform operations like computing the max attention value for each destination node.

1 Like

This helped a lot! thanks!!