to_captum: Support for explaining heterogenous graphs

[wsad1]

This PR adds
CaptumHeteroModel , to_captum_input and captum_output_to_dicts which can be used to explain heterogenous graphs as follows.

data: HeteroData = (...)
model = ... # A heterogenous model
mask_type = ...

captum_model: CaptumHeteroModel = to_captum_model(model, mask_type, output_idx=output_idx, data.metadata)
inputs, additonal_forward_args = to_captum_input(data.x_dict, data.edge_index_dict, mask_type)
ig = IntegratedGradients(captum_model)
ig_attr_nodes_edges = ig.attribute(inputs, target=target,
    additional_forward_args=additonal_forward_args, internal_batch_size=1)
x_attr_dict, edge_attr_dict = captum_output_to_dicts(ig_attr_nodes_edges, mask_type, data.metadata)

TODOs in follow up PRs

1. Add an example for to_captum with hetero data.
2. Move this behind the ExplainerAlgorithm interface being developed in GNN explanation settings  #5804. The interaace has to be extended to support HeteroData.

[codecov]

Codecov Report

Merging #5934 (b342bcc) into master (01037db) will increase coverage by 0.06%.
The diff coverage is 98.09%.

@@            Coverage Diff             @@
##           master    #5934      +/-   ##
==========================================
+ Coverage   84.70%   84.77%   +0.06%     
==========================================
  Files         360      360              
  Lines       19862    19962     +100     
==========================================
+ Hits        16825    16923      +98     
- Misses       3037     3039       +2     

Impacted Files	Coverage Δ
torch_geometric/nn/models/explainer.py	96.83% <98.07%> (+0.75%)	⬆️
torch_geometric/nn/models/__init__.py	100.00% <100.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[RBendias]

Thanks for the implementation! :D I have two points/suggestions:

1. I think currently the edge_mask is still the regular edge_index, but it needs to be a mask, i.e. torch.ones((1, num_edges )). This needs to be adjusted throughout the file.

2. Also, I am still in favor of a uniform implementation for hetero and non-hetero data. This way the functionality for the user would be the same regardless of the data type. Suggesting a method that converts the data and the model:

   • to_captum(data=None, model=None, ...) -> captum_input, captum_additional_args, CaptumModel / CaptumHeteroModel.

   or divided into two functions:

   • to_captum_data(...) -> captum_input, captum_additional_args

   • to_captum_model(...) -> CaptumModel / CaptumHeteroModel

   However, I would keep the input of these methods the same for hetero and non-hetero data and internally check the type of data and model.

[wsad1]

@RBendias added a method to_captum_input which returns the input and additonal_forward_args as you suggested. Will add tests in a while.

[RBendias]

LGTM! I left a few suggestions regardig the docs.