Return_attention_weights when using HeteroConv()

[Jerome-Cong]

Attention-based graph operator like GATv2Conv and TransformerConv have an option return_attention_weights when forward the layer. When using HeteroConv() to set up different aggregation for different edge types, for each layer different aggregation operation forward together like:

for conv in self.local_convs:
     h_dict = conv(h_dict, edge_index_dict)
     h_dict = {key: h.relu() for key, h in h_dict.items()}

Is there any way to pass the parameter return_attention_weights to the GATv2Conv operator for example?

[rusty1s]

Mh, this is an interesting thought. I guess this is not really supported at the moment. One idea is to simply create your own instance of HeteroConv:

outs = defaultdict(list)
for edge_type, edge_index in edge_index_dict.items():
    src, _, dst = edge_type
    out, attn = conv_dict[edge_type]((x_dict[src], x_dict[dst]), edge_index,
                                      return_attention_weights=True)
    outs[dst].append(out)

[Jerome-Cong]

Sorry for the late reply! I guess another simple idea is to break the HeteroConv apart (since I have two kinds of edges only). Thus if there's only one GATv2 layer in the HeteroConv dict, then I can do something like the follows:

for conv in convs:
     out, atten_weights = conv(h_dict, edge_index, return attention_weights=True)
     x_dict.update(out)

Is that possible?

[rusty1s]

Yes, this works as well, but I assume it would be something like

convs = ModuleDict({edge_type: GATConv(...)})

for edge_type, conv in convs.items():
    src, _, dst = edge_type
    out, weights = conv((h_dict[src], h_dict[dst]), edge_index_dict[edge_type], return_attention_weights=True)