Multilabel node classification explanation

[Laolga]

I'm trying to understand predictions of a multiclass and multilabel GNN that looks like this:

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(data.num_node_features, 16)
        self.conv_class1 = GCNConv(16, data.num_classes_class1 - 1)
        self.conv_class2 = GCNConv(16, data.num_classes_class2)


    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x_pre = F.dropout(x, training=self.training)
        x_class1 = self.conv_class1(x_pre, edge_index)
        x_class2= self.conv_class2(x_pre, edge_index)


        return (F.log_softmax(x_class1, dim=1),F.log_softmax(x_class2, dim=1))

And I tried to explain it like this:

explainer = Explainer(
    model=model,
    algorithm=GNNExplainer(epochs=200),
    explainer_config =dict(
        explanation_type='model',
        node_mask_type='attributes',
        edge_mask_type='object'),
    model_config=dict(
        mode='classification',
        task_level='node',
        return_type='raw',  
    ),
)

# Generate explanation for the node at index `10`:
explanation = explainer(data.x, data.edge_index, index=10)

And obviously failed because my model returns a tuple:

AttributeError                            Traceback (most recent call last)
Input In [146], in <cell line: 18>()
      1 explainer = Explainer(
      2     model=model,
      3     algorithm=GNNExplainer(epochs=200),
   (...)
     14     ),
     15 )
     17 # Generate explanation for the node at index `10`:
---> 18 explanation = explainer(data.x, data.edge_index, index=10)
     19 print(explanation.edge_mask)
     20 print(explanation.node_mask)

File /omics/groups/OE0540/internal/B260/users/olga/.conda/envs/torch_geo/lib/python3.8/site-packages/torch_geometric/explain/explainer.py:124, in Explainer.__call__(self, x, edge_index, target, index, target_index, **kwargs)
    120     if target is not None:
    121         warnings.warn(
    122             f"The 'target' should not be provided for the explanation "
    123             f"type '{explanation_type.value}'")
--> 124     target = self.get_prediction(x=x, edge_index=edge_index, **kwargs)
    126 training = self.model.training
    127 self.model.eval()

File /omics/groups/OE0540/internal/B260/users/olga/.conda/envs/torch_geo/lib/python3.8/site-packages/torch/autograd/grad_mode.py:27, in _DecoratorContextManager.__call__.<locals>.decorate_context(*args, **kwargs)
     24 @functools.wraps(func)
     25 def decorate_context(*args, **kwargs):
     26     with self.clone():
---> 27         return func(*args, **kwargs)

File /omics/groups/OE0540/internal/B260/users/olga/.conda/envs/torch_geo/lib/python3.8/site-packages/torch_geometric/explain/explainer.py:68, in Explainer.get_prediction(self, *args, **kwargs)
     66     out = self.model(*args, **kwargs)
     67 if self.model_config.mode == ModelMode.classification:
---> 68     out = out.argmax(dim=-1)
     70 self.model.train(training)
     72 return out

AttributeError: 'tuple' object has no attribute 'argmax'

What can I change to make it work?

Thanks!

[wsad1]

The explain module doesn't support multi-label classifications models at the moment. Instead you could have it explain just one of the labels with mode=multiclass_classification

[Laolga]

Do you mean I should change the model to a simple multiclass prediction?
Also, I guess my pytorch geometric is older since mode=multiclass_classification is not a valid ModelMode

[rusty1s]

Yes, multi_class_classification and binary_classification are only available on PyG master right now.

[wsad1]

Yeah explain can't handle mutilabel classification, so you'd have to change it to multiclass prediction.