Error obtaining attention weights from GATConv for heterogeneous graph

[enter802]

I am trying to get the attention weights of my GATConv network. However, my current procedure is failing to retrieve the weights and throwing out an error that I do not understand. I was able to create a minimalistic example to recreate the same error:

from torch_geometric.data import HeteroData
from torch_geometric.nn import to_hetero, GATConv
import torch

data = HeteroData()
data['type1'].x = torch.rand(3,2)
data['type2'].x = torch.rand(3,2)
data['type3'].x = torch.rand(3,2)

data['type1','to','type2'].edge_index = torch.tensor([[0, 1, 2],
                                                      [1, 1, 2]]).long().contiguous()

data['type1','to','type2'].edge_attr = torch.rand(3,1)

data['type2','to','type1'].edge_index = torch.tensor([[0, 1, 2],
                                                      [2, 1, 1]]).long().contiguous()

data['type2','to','type1'].edge_attr = torch.rand(3,1)

data['type1','to','type3'].edge_index = torch.tensor([[0, 1],
                                                      [1, 1]]).long().contiguous()

data['type1','to','type3'].edge_attr = torch.rand(2,1)

data['type2','to','type3'].edge_index = torch.tensor([[0, 1, 2], # Comment this to stop the error
                                                      [0, 1, 2]]).long().contiguous()

data['type2','to','type3'].edge_attr = torch.rand(3,1) # Comment this to stop the error

class GAT_Individual(torch.nn.Module):
    def __init__(self, output_dim, seed):
        super().__init__()
        torch.manual_seed(seed)
        self.GNN = GATConv((-1, -1), output_dim, add_self_loops=False)
    def forward(self, x, edge_index, edge_attr):
        x, alpha = self.GNN(x, edge_index, edge_attr, return_attention_weights=True)
        return x, alpha

model = GAT_Individual(2,10)

model = to_hetero(model, data.metadata(), aggr='sum')

x, alpha = model(data.x_dict, data.edge_index_dict, data.edge_attr_dict)
print(alpha)

The above code produces the following error:

TypeError: add(): argument 'input' (position 1) must be Tensor, not tuple

However, when I comment out the edge connections and edge attributes between "type2" and "type3" nodes (indicated above in the code), the error disappears. I am not sure what's happening. Your help will be highly appreciated. Thanks in advance!

[rusty1s]

I am sorry to inform you that return_attention_weights=True and to_hetero is currently not compatible, and I don't see a good way to support this long-term. The only workaround would be to implement manually what to_hetero is doing internally, i.e., specifying a GATConv layer for every edge type, and then calling these layers in a loop:

self.convs = ModuleDict({edge_type: GATConv((-1, -1), output_dim, ...) for edge_type in edge_types})

def forward(self, x_dict, edge_index_dict):
    for edge_type, edge_index in edge_index_dict.items():
         src, _, dst = edge_type
         
         x, alpha = self.convs[edge_type]((x_dict[src], x_dict[dst]), edge_index, return_attention_weights=True)

[enter802]

Thanks a lot. I understand the difficulty of processing a heterogeneous graph. Let me try to workaround with your suggesting!

Updated:
IMO: It is important to have this feature since it helps in explaining the algorithm.

[xubingze]

Hi @enter802 , I wonder have you solved this problem? I am trying to do similar things and would like your advice on this.

[enter802]

@xubingze this worked for me:

from torch_geometric.data import HeteroData
from torch_geometric.nn import to_hetero, GATConv
import torch

data = HeteroData()
data['type1'].x = torch.rand(5,2)
data['type2'].x = torch.rand(3,2)
data['type3'].x = torch.rand(3,2)

data['type1','to','type2'].edge_index = torch.tensor([[0, 1, 2, 4],
                                                      [1, 0, 1, 0]]).long().contiguous()

data['type1','to','type2'].edge_attr = torch.rand(3,1)

data['type2','to','type1'].edge_index = torch.tensor([[0, 1, 2],
                                                      [2, 1, 1]]).long().contiguous()

data['type2','to','type1'].edge_attr = torch.rand(3,1)

data['type1','to','type3'].edge_index = torch.tensor([[0, 1],
                                                      [1, 1]]).long().contiguous()

data['type1','to','type3'].edge_attr = torch.rand(2,1)

data['type2','to','type3'].edge_index = torch.tensor([[0, 1, 2],
                                                      [0, 1, 2]]).long().contiguous()

data['type2','to','type3'].edge_attr = torch.rand(3,1)


class GAT_Individual(torch.nn.Module):
    def __init__(self, output_dim, seed):
        super().__init__()
        torch.manual_seed(seed)
        self.GNN = GATConv((-1, -1), output_dim, add_self_loops=False)
    def forward(self, x, edge_index, edge_attr):
        x = self.GNN(x, edge_index, edge_attr)
        return x
    
def attnWeights(model, x_dict, edge_index_dict, edge_attr_dict):
    for (edge_type, edge_index), (_, edge_attr) in zip(edge_index_dict.items(),edge_attr_dict.items()):
         src, _, dst = edge_type
         result = '__'.join(edge_type)
         x, alpha = model.GNN[result]((x_dict[src], x_dict[dst]), edge_index, edge_attr, 
                                      return_attention_weights=True)
         
         print(alpha) # edges and their corresponding attention values
         
    return x, alpha

model = GAT_Individual(2,10)

model = to_hetero(model, data.metadata(), aggr='sum')

_ = model(data.x_dict, data.edge_index_dict, data.edge_attr_dict)

x, alpha = attnWeights(model, data.x_dict, data.edge_index_dict, data.edge_attr_dict)

[xubingze]

@enter802 Thank you so much for sharing! I am wondering what is "ModuleDict" in rusty's answer, and your examples exactly solve my problems. Let me try this on my own datasets. I really appreciate your help and timely response!

[xubingze]

@enter802 So if I would like to train the model with two GATConv layers and then get the attention weights in the last layer. What shall I do? Are the following steps correct?
1, Follow the common model.train() steps in training the model. Since I have two GATConv, in the forward functions I would add a self.GNN2(x, edge_index, edge_attr) layer.

class GAT_Individual(torch.nn.Module):
    def __init__(self, output_dim, seed):
        super().__init__()
        torch.manual_seed(seed)
        self.GNN = GATConv((-1, -1), output_dim, add_self_loops=False)
        self.GNN2 = GATConv((-1, -1), output_dim, add_self_loops=False)
    def forward(self, x, edge_index, edge_attr):
        x = self.GNN(x, edge_index, edge_attr).relu()
        x = self.GNN2(x, edge_index, edge_attr)
        return x

2, Then the attnWeights function should be called after "@torch.no_grad()" similar to the testing steps. In this case, to get the weights in the last layers, (I am not too sure about this step)

def attnWeights(model, x_dict, edge_index_dict, edge_attr_dict):
    for (edge_type, edge_index), (_, edge_attr) in zip(edge_index_dict.items(),edge_attr_dict.items()):
         src, _, dst = edge_type
         result = '__'.join(edge_type)
         x, alpha_1 = model.GNN[result]((x_dict[src], x_dict[dst]), edge_index, edge_attr, 
                                      return_attention_weights=True)
         x, alpha_2 = model.GNN2[result]((x_dict[src], x_dict[dst]), edge_index, edge_attr, 
                                      return_attention_weights=True)         
    return x, alpha_2

, where alpha_2 is the attention weights in the last layers.

3, Do the results (x) get from attnWeights align with the results (x) from model(data.x_dict, data.edge_index_dict, data.edge_attr_dict)?

Thank you very much for your feedback.