Explainer algorithms for heterogeneous graph.

[viplovearora]

Hello, I am trying to use the pyg explain module for a heterogeneous graph in the link prediction setting. Although explainer algorithms like GNNExplainer and AttentionExplainer allow heterogeneous graphs as input, when you try to run them, an error saying "Heterogeneous graphs not yet supported" pops up. Has this been implemented already or am I doing something wrong?

I tried using CaptumExplainer as well but I am getting the following error: assert dims == 2, "Output must be 2D to select tensor of targets." I am looking into this too, but any help will be much appreciated. Thanks!

[rusty1s]

Currently, only CaptumExplainer supports heterogeneous graphs. Do you have a small example to show your error? Which output shape does your model return?

[viplovearora]

Thank you for the prompt response @rusty1s. Will the other explainer models be available for heterogeneous graphs in the future?

My model returns a single tensor as shown below:

class GNNEncoder(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.conv1 = GATv2Conv((-1, -1), hidden_channels, heads=1, add_self_loops=False)
        self.conv2 = GATv2Conv((-1, -1), hidden_channels, heads=1, add_self_loops=False)

    def forward(self, x, edge_index, eval=False):
        x, (indices,attention_score) = self.conv1(x, edge_index, return_attention_weights=True)
        x.relu()
        x, (indices,attention_score) = self.conv2(x, edge_index, return_attention_weights=True)
        return x
 

class EdgeDecoder(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.lin1 = Linear(2 * hidden_channels, hidden_channels)
        self.lin2 = Linear(hidden_channels, 1)

    def forward(self, z_dict, edge_label_index):
        row, col = edge_label_index
        z = torch.cat([z_dict['protein'][row], z_dict['transcript'][col]], dim=-1)

        z = self.lin1(z).relu()
        z = self.lin2(z)
        return z.view(-1)


class Model(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.encoder = GNNEncoder(hidden_channels)
        self.encoder = to_hetero(self.encoder, data.metadata(), aggr='sum')
        self.decoder = EdgeDecoder(hidden_channels)

    def forward(self, x_dict, edge_index_dict, edge_label_index):
        z_dict = self.encoder(x_dict, edge_index_dict)
        return self.decoder(z_dict, edge_label_index)

I am using this example for using CaptumExplainer with my dataset. And I get the following error:

AssertionError                            Traceback (most recent call last)
Cell In[19], line 1
----> 1 explanation = explainer(
      2     data.x_dict,
      3     data.edge_index_dict,
      4     index=index,
      5     edge_label_index=torch.cat([data['protein', 'transcript'].edge_index, data['protein', 'transcript'].train_neg_edge_index], dim=1),
      6     # target=torch.cat([torch.ones(num_train), torch.zeros(num_train)]).to(device),
      7 )
      8 print(f'Generated explanations in {explanation.available_explanations}')

File [~/miniconda3/lib/python3.9/site-packages/torch_geometric/explain/explainer.py:198](https://vscode-remote+ssh-002dremote-002btimavo.vscode-resource.vscode-cdn.net/u/v/varora/Desktop/code/RPI_salmon/~/miniconda3/lib/python3.9/site-packages/torch_geometric/explain/explainer.py:198), in Explainer.__call__(self, x, edge_index, target, index, **kwargs)
    195 training = self.model.training
    196 self.model.eval()
--> 198 explanation = self.algorithm(
    199     self.model,
    200     x,
    201     edge_index,
    202     target=target,
    203     index=index,
    204     **kwargs,
    205 )
    207 self.model.train(training)
    209 # Add explainer objectives to the `Explanation` object:
...
--> 503     assert dims == 2, "Output must be 2D to select tensor of targets."
    504     return torch.gather(output, 1, target.reshape(len(output), 1))
    505 else:

AssertionError: Output must be 2D to select tensor of targets.

I have seen a similar issue in the Captum issues pytorch/captum#288 but I don't quite understand how to solve this within the pyg call.

Can you also explain how I can modify my model to get the attention scores?

[rusty1s]

Oh, this is related to binary classification with CaptumExplainer? We just fixed this a few weeks ago on master. Can you update to master and see if that fixes your issues?

[viplovearora]

Yes, it is the binary classification, and it works after the update. Thanks for the help @rusty1s.
Can you also explain how I can modify my model to get the attention scores?