Call visualize_graph on pytorch_geometric version 2.2.0

[PeeteKeesel]

Hey all 👋🏽
I have build a bi-modal regression GNN and now would like to run GNNExplainer on it. Let me explain my problem.
✅ What works: I was able to successfully initialize the GNNExplainer and call explain_graph.

explainer = GNNExplainer(
    model=model.model,
    epochs=100,
    return_type='regression',
    nhops=2
)

node_feat_mask, edge_mask = explainer.explain_graph(
    x=G.x, 
    edge_index=G.edge_index,
    cell_batch=torch.tensor([0]),
    drug=torch.tensor(drug).double().reshape(1, 256)
)

Now I wanted to call visualize_graph on the explainer but saw that this method got deleted after version 2.0.3.

So, currently it is not possible to visualize a (sub-)graph on an initialized GNNExplainer. Not really sure why that is but I believe it has it's reasons.

⛔ What does not work: Now I saw in the provided gnn_explainer.py example file that it should in the newest version be called like that:

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

That works. However when I now want to call

explanation = explainer(
    x=G.x,
    edge_index=G.edge_index,
    cell_batch=torch.tensor([0]),
    drug=torch.tensor(drug).double().reshape(1, 256)
)

it throws the error that

TypeError: MyGNN.forward() got an unexpected keyword argument 'x'

which is makes sense since my forward method looks like this:

    def forward(self, cell_x, cell_edge_index, cell_batch, drug):
        drug_emb = self.drug_emb(drug)
        cell_emb = self.cell_emb(cell_x, cell_edge_index, cell_batch)
    # etc.

However, I don't want to change the input parameter names of my forward method to x and edge_index. The reason is that I will add two different x and edge_index in future iteration such that I have a bi-modal network with two GNNs. So that's why I can't change it to x and edge_index since in the future there will be basically x2 and edge_index2. So I would like to keep my namings.

🤔 Is there any other way how I can run explainer while keeping the naming of the arguments of my forward method?

[rusty1s]

You should be able to call it via

explanation = explainer(
    G.x,
    G.edge_index,
    cell_batch=torch.tensor([0]),
    drug=torch.tensor(drug).double().reshape(1, 256)
)

[rusty1s]

See https://github.com/pyg-team/pytorch_geometric#nightly-and-master

[PeeteKeesel]

Thanks!!! I tried that but got still the same error as above.

[PeeteKeesel]

FYI:

The full error message is as follows:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [20], in <cell line: 1>()
----> 1 explanation = explainer(
      2     G.x,
      3     G.edge_index,
      4     cell_batch=torch.tensor([0]),
      5     drug=torch.tensor(drug).double().reshape(1, 256)
      6 )

File ~/anaconda3/envs/myproject/lib/python3.10/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 ~/anaconda3/envs/myproject/lib/python3.10/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 ~/anaconda3/envs/myproject/lib/python3.10/site-packages/torch_geometric/explain/explainer.py:66, in Explainer.get_prediction(self, *args, **kwargs)
     63 self.model.eval()
     65 with torch.no_grad():
---> 66     out = self.model(*args, **kwargs)
     67 if self.model_config.mode == ModelMode.classification:
     68     out = out.argmax(dim=-1)

File ~/anaconda3/envs/myproject/lib/python3.10/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

TypeError: MyGNN.forward() got an unexpected keyword argument 'x'

after I ran

!pip install git+https://github.com/pyg-team/pytorch_geometric.git

from torch_geometric.explain import Explainer, GNNExplainer

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

explanation = explainer(
    G.x,
    G.edge_index,
    cell_batch=torch.tensor([0]),
    drug=torch.tensor(drug).double().reshape(1, 256)
)

I am using

torch                     1.12.1                   pypi_0    pypi
torch-cluster             1.6.0                    pypi_0    pypi
torch-geometric           2.2.0                    pypi_0    pypi
torch-scatter             2.1.0                    pypi_0    pypi
torch-sparse              0.6.15                   pypi_0    pypi
torchaudio                0.12.1                   pypi_0    pypi
torchvision               0.13.1                   pypi_0    pypi
python                    3.10.4               h12debd9_0  

[rusty1s]

I just tried this myself by renaming the model arguments in examples/gnn_explainer.py and it works just fine. Are you sure you are on latest master? You may need to enforce a re-installation:

pip install --force-reinstall git+https://github.com/pyg-team/pytorch_geometric.git

[PeeteKeesel]

Interesting. Hmm, I tried with the --force-reinstall but now getting

ERROR: Could not install packages due to an OSError: [Errno 28] No space left on device

I will try to solve this first.

Thanks for running by yourself for checking 🙏🏽