Explain module with hetero graph and edges prediction, not working

[arthurserres]

Hello, I am interested in a specific case of the explain library which I cannot manage to solve.

More precisely, I have a model on HeteroData Graph that is performing edge classification. In particular, my model output is not a unique tensor but rather a 'dictionary' of tensors. (One tensor for each type of data). On this example It is mentioned that 'It is assumed that model outputs a single tensor' which is not the case in my example. Indeed, if I try to call the explainer I obtain the error 'dict' object has no attribute 'argmax' ...

I would like to know how I can, and I should use the explain library with my settings ? (I already checked the examples page but it does not help)

[akihironitta]

@arthurserres Could I see your script? What does your model look like?

I feel like the example here will still be a reasonable reference for your use-case.

[arthurserres]

Hi akihironitta, thanks a lot for the answer. My graph has two different types of edges, let us say type1 and type2. Therefore, if we compute model(batch.x_dict, batch.edge_index_dict, batch.edge_attr_dict), we obtain a dictionary of tensors where each element of the dictionary corresponds to the prediction of each edge type type1 and type2. I can provide the exact architecture of the model, but I think that the 'issue' here is that instead of a tensor the output is a dictionary of tensors ( I understood it is quite standard with HeteroData).

For now, the solution I found is to consider a new class of model:

class model_trunc(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    def forward(self, x_dict, edge_index_dict, edge_attr_dict):
        return self.model(x_dict,edge_index_dict, edge_attr_dict)[('type1')]

and to perform the explanation on Model = model_trunc(model). I think it does the job as I am only interested in the explanations of type1 and here the output of Model is a single tensor.

Looking back at the example that you provide, I think this is essentially what they do with this EdgeDecoder (They are making sure the output is a single tensor).

Thanks a lot.

Best,

Arthur

[rusty1s]

Yes, outputs needs to be a torch.Tensor. Your solution works. Alternatively, you can stack the outputs of the two node types together before returning it.