Problem with trying to explain heterogeneous model built using to_hetero

[bl3e967]

Hello, I am relatively new to pytorch and pyg and had some questions with regards to the explainability functionality.

I have the following heterogeneous dataset:

HeteroData(
  internal={
    x=[4071037, 84],
    y=[4071037],
    train_mask=[2068189],
    test_mask=[2002848]
  },
  external={ x=[20410080, 74] },
  (internal, connection, internal)={
    edge_index=[2, 2070823],
    edge_attr=[2070823, 50]
  },
  (internal, connection, external)={
    edge_index=[2, 13740742],
    edge_attr=[13740742, 46]
  },
  (internal, connection, internal)={
    edge_index=[2, 19553334],
    edge_attr=[19553334, 54]
  }
)

and I am building a GATConv model by using the default GAT model in torch_geometric.nn.models.GAT and the to_hetero method:

model = GAT(*args, **kwargs)
model = to_hetero(model, metadata)

to get the below model:

GraphModule(
  (_trim): ModuleDict(
    (internal__connection__internal): TrimToLayer()
    (internal__connection__external): TrimToLayer()
    (external__connection__internal): TrimToLayer()
  )
  (convs): ModuleList(
    (0): ModuleDict(
      (internal__connection__internal): GATConv((-1, -1), 51, heads=3)
      (internal__connection__external): GATConv((-1, -1), 51, heads=3)
      (external__connection__internal): GATConv((-1, -1), 51, heads=3)
    )
    (1-2): 2 x ModuleDict(
      (internal__connection__internal): GATConv((153, 153), 51, heads=3)
      (internal__connection__external): GATConv((153, 153), 51, heads=3)
      (external__connection__internal): GATConv((153, 153), 51, heads=3)
    )
  )
  (norms): ModuleList(
    (0-2): 3 x ModuleDict(
      (internal): BatchNorm(153)
      (external): BatchNorm(153)
    )
  )
  (act): ModuleDict(
    (internal): ReLU()
    (external): ReLU()
  )
  (jk): ModuleDict(
    (internal): JumpingKnowledge(max)
    (external): JumpingKnowledge(max)
  )
  (lin): ModuleDict(
    (internal): Linear(in_features=153, out_features=1, bias=True)
    (external): Linear(in_features=153, out_features=1, bias=True)
  )
)

The model outputs a dictionary of tensors keyed by the node types, which is expected behaviour according to the docs.

However, the problem arises when I try and run this model through the Explainer, as the explainer expects a heterogeneous model to output a single tensor, not a dictionary of tensors.

For example, using Captum integrated gradients

explainer = Explainer(
    model=model,
    algorithm=CaptumExplainer('IntegratedGradients'),
    explanation_type='model',
    node_mask_type='attributes',
    edge_mask_type='object',
    model_config=dict(
        mode='binary_classification',
        task_level="node",
        return_type="raw"
    )
)

gives the following error:

AttributeError                            Traceback (most recent call last)
Cell In[23], line 1
----> 1 train_explanation = explainer(
      2     x=batch.x_dict, 
      3     edge_index=batch.edge_index_dict, 
      4     edge_attr=batch.edge_attr_dict,
      5 )

File /opt/conda/envs/gnn/lib/python3.10/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:

File /opt/conda/envs/gnn/lib/python3.10/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []

File /opt/conda/envs/gnn/lib/python3.10/site-packages/torch_geometric/explain/algorithm/captum_explainer.py:162, in CaptumExplainer.forward(self, model, x, edge_index, target, index, **kwargs)
    159 else:
    160     target = target[index]
--> 162 attributions = self.attribution_method.attribute(
    163     inputs=inputs,
    164     target=target,
    165     additional_forward_args=add_forward_args,
    166     **self.kwargs,
    167 )
    169 node_mask, edge_mask = convert_captum_output(
    170     attributions,
    171     mask_type,
    172     metadata,
    173 )
    175 if not isinstance(x, dict):

File /opt/conda/envs/gnn/lib/python3.10/site-packages/captum/log/__init__.py:42, in log_usage.<locals>._log_usage.<locals>.wrapper(*args, **kwargs)
     40 @wraps(func)
     41 def wrapper(*args, **kwargs):
---> 42     return func(*args, **kwargs)

File /opt/conda/envs/gnn/lib/python3.10/site-packages/captum/attr/_core/integrated_gradients.py:274, in IntegratedGradients.attribute(self, inputs, baselines, target, additional_forward_args, n_steps, method, internal_batch_size, return_convergence_delta)
...
--> 494 num_examples = output.shape[0]
    495 dims = len(output.shape)
    496 device = output.device

AttributeError: 'dict' object has no attribute 'shape'

Could someone please advise if I am doing something wrong with the Explainer, or if this is expected behaviour? If this is expected, is there any way I can use the to_hetero method while getting the explainability functionality working?

[rusty1s]

I think the issue is that your heterogeneous model is still expected to return single tensors, not dictionaries of tensors. As such, what you can do is something like

class HeteroSAGE(torch.nn.Module):
    def __init__(self, metadata):
        super().__init__()
        self.graph_sage = to_hetero(GraphSAGE(-1, 32, num_layers=2), metadata)
        self.lin = torch.nn.Linear(32, 1)

    def forward(self, x_dict, edge_index_dict) -> torch.Tensor:
        return self.lin(self.graph_sage(x_dict, edge_index_dict)['paper'])

[bl3e967]

Yep that's done it thanks.