GNNExplainer on AttentiveFP

[Takaogahara]

I'm trying to use the GNNExplainer on AttentiveFP, but I'm having some difficulties and I can't say for sure if the problem is me or if I should open an issue.

I get a TypeError error: forward() got multiple values for argument 'edge_attr' when passing the edge_attr **kwarg in the explain_graph function.
I used GNNExplainer for GAT and GCN and had no problems

I adapted the example available from GNNExplainer to reproduce the described error.

import os.path as osp
import torch
import torch.nn as nn
from torch_geometric.datasets import MoleculeNet
from torch_geometric.nn import AttentiveFP, GNNExplainer

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


dataset = 'bbbp'
path = osp.join(osp.dirname(osp.realpath(__file__)), 'MoleculeNet')
dataset = MoleculeNet(path, dataset)
data = dataset[0]


net = AttentiveFP(in_channels=data.num_node_features,
                  hidden_channels=64,
                  out_channels=1,
                  edge_dim=data.num_edge_features,
                  num_layers=1,
                  num_timesteps=1,
                  dropout=0.2)


model = net.to(device)
data = data.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
data_batch = torch.zeros(data.x.shape[0], dtype=int, device=device)
x = x.float()

for epoch in range(1, 201):
    model.train()
    optimizer.zero_grad()
    logit = model(x=x, edge_index=edge_index,
                  edge_attr=edge_attr, batch=data_batch)

    loss_fn = nn.BCEWithLogitsLoss()
    loss = loss_fn(logit, data.y)
    loss.backward()
    optimizer.step()

explainer = GNNExplainer(model, epochs=200, return_type='log_prob')
_, edge_mask = explainer.explain_graph(x=x, edge_index=edge_index,
                                       edge_attr=edge_attr)

print(edge_mask)

This is the Traceback I received.

Traceback (most recent call last):
  File "/media/takaogahara/data/projects/gnn-toolkit/temp/issue.py", line 45, in <module>
    _, edge_mask = explainer.explain_graph(x=x, edge_index=edge_index,
  File "/home/takaogahara/virtualenvs/gnn-toolkit/lib/python3.8/site-packages/torch_geometric/nn/models/gnn_explainer.py", line 146, in explain_graph
    prediction = self.get_initial_prediction(x, edge_index, batch=batch,
  File "/home/takaogahara/virtualenvs/gnn-toolkit/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/takaogahara/virtualenvs/gnn-toolkit/lib/python3.8/site-packages/torch_geometric/nn/models/explainer.py", line 235, in get_initial_prediction
    out = self.model(x, edge_index, batch, **kwargs)
  File "/home/takaogahara/virtualenvs/gnn-toolkit/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
TypeError: forward() got multiple values for argument 'edge_attr'

Has anyone used GNNExplainer in AttentiveFP and could help me?

[wsad1]

This looks like a bug in Explainer.get_prediction will try to send in a fix soon.
The problem is that in attentive_fp the batch argument comes after the edge_attr.

[wsad1]

Should be fixed now. try installing from master with pip install git+https://github.com/pyg-team/pytorch_geometric.git.

[Takaogahara]

Thanks !!