Repository: pyg-team/pytorch_geometric

Issue:
CaptumExplainer binary graph classification
### 🐛 Describe the bug

Hello,

A while ago, I commented in #7702 about the impossibility of using CaptumExplainer with binary classification, which was corrected in PR #7787. However, I only use the fixed implementation now due to the release update.

However, a function converts binary classification to multi-class classification in `torch_geometric.explain.algorithm.captum.py` postprocess function (line 88).

Here is a code to reproduce the error:

```python
import torch
from tqdm import tqdm
import os.path as osp
import torch.nn as nn
from torch_geometric.loader import DataLoader
from torch_geometric.datasets import MoleculeNet
from torch_geometric.nn import GAT, global_mean_pool
from torch_geometric.explain import Explainer, CaptumExplainer


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


dataset = "hiv"
path = osp.join(osp.dirname(osp.realpath(__file__)), "MoleculeNet")
dataset = MoleculeNet(path, dataset)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

model = GAT(in_channels=dataset.num_node_features,
            hidden_channels=64,
            out_channels=1,
            num_layers=1,
            dropout=0.2,
            edge_dim=dataset.num_edge_features,
            v2=True,
            heads=2)

model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.BCEWithLogitsLoss()

model.train()
for epoch in tqdm(range(1, 2), ncols=120):
    for batch in loader:
        batch.to(device)
        optimizer.zero_grad()

        out = model(x=batch.x.float(), edge_index=batch.edge_index,
                    edge_attr=batch.edge_attr.float())
        out = global_mean_pool(out, batch.batch)
        pred = torch.sigmoid(out)

        loss = loss_fn(pred, batch.y)
        loss.backward()
        optimizer.step()


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="graph",
                                        return_type="probs"))

data = dataset[0].to(device)
dummy = torch.zeros(data.x.shape[0], dtype=int, device=device)

explanation = explainer(x=data.x.float(), edge_index=data.edge_index,
                        edge_attr=data.edge_attr.float())

print(explanation.available_explanations)
```

This is the Traceback I received.

``` bash
Traceback (most recent call last):
  File "/home/takaogahara/takaogahara/experiments/explain/issue.py", line 60, in <module>
    explanation = explainer(x=data.x.float(), edge_index=data.edge_index,
  File "/home/takaogahara/virtualenvs/graphchem/lib/python3.10/site-packages/torch_geometric/explain/explainer.py", line 204, in __call__
    explanation = self.algorithm(
  File "/home/takaogahara/virtualenvs/graphchem/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/takaogahara/virtualenvs/graphchem/lib/python3.10/site-packages/torch_geometric/explain/algorithm/captum_explainer.py", line 167, in forward
    attributions = self.attribution_method_instance.attribute(
  File "/home/takaogahara/virtualenvs/graphchem/lib/python3.10/site-packages/captum/log/__init__.py", line 42, in wrapper
    return func(*args, **kwargs)
  File "/home/takaogahara/virtualenvs/graphchem/lib/python3.10/site-packages/captum/attr/_core/integrated_gradients.py", line 274, in attribute
    attributions = _batch_attribution(
  File "/home/takaogahara/virtualenvs/graphchem/lib/python3.10/site-packages/captum/attr/_utils/batching.py", line 78, in _batch_attribution
    current_attr = attr_method._attribute(
  File "/home/takaogahara/virtualenvs/graphchem/lib/python3.10/site-packages/captum/attr/_core/integrated_gradients.py", line 351, in _attribute
    grads = self.gradient_func(
  File "/home/takaogahara/virtualenvs/graphchem/lib/python3.10/site-packages/captum/_utils/gradient.py", line 112, in compute_gradients
    outputs = _run_forward(forward_fn, inputs, target_ind, additional_forward_args)
  File "/home/takaogahara/virtualenvs/graphchem/lib/python3.10/site-packages/captum/_utils/common.py", line 487, in _run_forward
    return _select_targets(output, target)
  File "/home/takaogahara/virtualenvs/graphchem/lib/python3.10/site-packages/captum/_utils/common.py", line 506, in _select_targets
    raise AssertionError(
AssertionError: Tensor target dimension torch.Size([1, 19]) is not valid. torch.Size([19, 2])
```

I believe the postprocess function in the `captum.py` file needs to be updated to keep the binary classification data, am I right?

### Environment

* PyG version: 2.4.0
* PyTorch version: 2.0.1+cu118
* OS: Ubuntu 22.04
* Python version: 3.10
* CUDA/cuDNN version: cu118
* How you installed PyTorch and PyG (`conda`, `pip`, source): pip


Discussion:
I believe I am having the same issue, though I'm not entirely certain what the expected output of the model should be for binary classification (or for any of the options, actually) to confirm that it's not an issue with my code as opposed to the model itself. My setup is nearly identical to the minimally reproducible example by @Takaogahara above with the exception of the graph and model in question being heterogenous.
The issue with the code above is that the global pooling is performed outside of the loop, thus making configuration and model not align anymore. This fixes it for me:
```python
import os.path as osp

import torch
import torch.nn as nn

from torch_geometric.datasets import MoleculeNet
from torch_geometric.explain import CaptumExplainer, Explainer
from torch_geometric.loader import DataLoader
from torch_geometric.nn import GAT, global_mean_pool

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

dataset = "hiv"
path = osp.join(osp.dirname(osp.realpath(__file__)), "MoleculeNet")
dataset = MoleculeNet(path, dataset)
loader = DataLoader(dataset, batch_size=32, shuffle=True)


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = GAT(
            in_channels=dataset.num_node_features,
            hidden_channels=64,
            out_channels=1,
            num_layers=1,
            dropout=0.2,
            edge_dim=dataset.num_edge_features,
            v2=True,
            heads=2,
        )

    def forward(self, x, edge_index, edge_attr, batch):
        x = self.model(x, edge_index, edge_attr=edge_attr)
        return global_mean_pool(x, batch)


model = Model().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.BCEWithLogitsLoss()

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="graph",
                      return_type="probs"),
)

data = dataset[0].to(device)

explanation = explainer(
    x=data.x.float(),
    edge_index=data.edge_index,
    edge_attr=data.edge_attr.float(),
    batch=torch.zeros(data.num_nodes, dtype=torch.long, device=device),
)

print(explanation.available_explanations)

Resolving patch:
diff --git a/torch_geometric/explain/algorithm/captum.py b/torch_geometric/explain/algorithm/captum.py
--- a/torch_geometric/explain/algorithm/captum.py
+++ b/torch_geometric/explain/algorithm/captum.py
@@ -236,7 +236,9 @@ def to_captum_input(
         raise ValueError(
             "'x' and 'edge_index' need to be either"
             f"'Dict' or 'Tensor' got({type(x)}, {type(edge_index)})")
+
     additional_forward_args.extend(args)
+
     return tuple(inputs), tuple(additional_forward_args)
 
 
diff --git a/torch_geometric/explain/algorithm/captum_explainer.py b/torch_geometric/explain/algorithm/captum_explainer.py
--- a/torch_geometric/explain/algorithm/captum_explainer.py
+++ b/torch_geometric/explain/algorithm/captum_explainer.py
@@ -138,7 +138,7 @@ def forward(
             *kwargs.values(),
         )
 
-        if isinstance(x, dict):
+        if isinstance(x, dict):  # Heterogeneous GNN:
             metadata = (list(x.keys()), list(edge_index.keys()))
             captum_model = CaptumHeteroModel(
                 model,
@@ -147,21 +147,24 @@ def forward(
                 metadata,
                 self.model_config,
             )
-        else:
+        else:  # Homogeneous GNN:
             metadata = None
-            captum_model = CaptumModel(model, mask_type, index,
-                                       self.model_config)
+            captum_model = CaptumModel(
+                model,
+                mask_type,
+                index,
+                self.model_config,
+            )
 
         self.attribution_method_instance = self.attribution_method_class(
             captum_model)
 
-        # In captum, the target is the class index for which
-        # the attribution is computed. With CaptumModel, we transform
-        # the binary classification into a multi-class. This way we can
-        # explain both classes and need to pass a target here as well.
+        # In Captum, the target is the class index for which the attribution is
+        # computed. Within CaptumModel, we transform the binary classification
+        # into a multi-class classification task.
         if self.model_config.mode == ModelMode.regression:
             target = None
-        else:
+        elif index is not None:
             target = target[index]
 
         attributions = self.attribution_method_instance.attribute(


Test patch: