How to explain the output of a Pyg-network (Graph conv) without target?

[monk1337]

Hi, I am trying to find ways to visualize(explain) the output of any pyg network. For example, I am taking a simple toy network using Graph conv:

import numpy as np
import torch

num_of_nodes = 4
num_of_feats = 5
output_dim   = 6

def get_data():
  adj         = torch.Tensor(np.random.randint(0,2,[num_of_nodes,num_of_nodes]))
  feats       = torch.Tensor(np.random.uniform(-1,1,[num_of_nodes, num_of_feats]))

  edge_indexs = adj.nonzero(as_tuple=False).t()
  return feats, edge_indexs

The network looks like this:

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(num_of_feats, output_dim)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        return x

calling network:

feats, ed_in = get_data()
xg = GCN()(feats, ed_in)
xg

output:

tensor([[0.4645, 0.4537, 0.7191, 0.2981, 0.9345, 0.2394],
        [0.4699, 0.3964, 0.4689, 0.5432, 0.4923, 0.5697],
        [0.0000, 0.0000, 0.0019, 0.0000, 0.1044, 0.0000],
        [0.5781, 0.5397, 0.6135, 0.5073, 0.7209, 0.5636]],
       grad_fn=<ReluBackward0>)

Is there a way to visualize this output?

I found two tutorials

1. Explaining GNN Model Predictions using Captum
2. GNNExplainer Tutorial

The tutorial uses molecules as the dataset, I am looking for a very simple example of how to visualize the output of the above network using captum?

What I have tried:

I tried to explain the embeddings using example given in repo

from torch_geometric.utils.convert import from_networkx
from torch_geometric.nn import Explainer, GCNConv, to_captum
from captum.attr import IntegratedGradients

import matplotlib.pyplot as plt

import numpy as np
import networkx as nx


def numpy_to_graph(A, node_features=None):
    
    G = nx.from_numpy_array(A)
    if node_features != None:
        for n in G.nodes():
            for k,v in node_features.items():
                G.nodes[n][k] = v[n]
    return G


def get_data(adj, feats):
    G           = numpy_to_graph(adj, node_features={'x': feats})
    pyg_graph   = from_networkx(G)
    return pyg_graph


def captum_explain_graph(model, 
                         edge_index, 
                         adj_matrix, 
                         output_idx, 
                         device):

  data   = get_data(adj_matrix, edge_index)
  target = 0
  model  = model.double()

  # Edge explainability
  # ===================

  # Captum assumes that for all given input tensors, dimension 0 is
  # equal to the number of samples. Therefore, we use unsqueeze(0).
  captum_model = to_captum(model, mask_type='edge', output_idx=output_idx)
  edge_mask    = torch.ones(data.num_edges, requires_grad=True, device=device)

  ig           = IntegratedGradients(captum_model)
  ig_attr_edge = ig.attribute(edge_mask.unsqueeze(0), target=target,
                              additional_forward_args=(data.x, data.edge_index),
                              internal_batch_size=1)

  # Scale attributions to [0, 1]:
  ig_attr_edge = ig_attr_edge.squeeze(0).abs()
  ig_attr_edge /= ig_attr_edge.max()

  # Visualize absolute values of attributions:
  explainer = Explainer(model)
  ax, G = explainer.visualize_subgraph(output_idx, data.edge_index, ig_attr_edge)
  plt.show()

  # Node explainability
  # ===================

  captum_model = to_captum(model, mask_type='node', output_idx=output_idx)

  ig           = IntegratedGradients(captum_model)
  ig_attr_node = ig.attribute(data.x.unsqueeze(0), target=target,
                              additional_forward_args=(data.edge_index),
                              internal_batch_size=1)

  # Scale attributions to [0, 1]:
  ig_attr_node = ig_attr_node.squeeze(0).abs().sum(dim=1)
  ig_attr_node /= ig_attr_node.max()

  # Visualize absolute values of attributions:
  ax, G = explainer.visualize_subgraph(output_idx, data.edge_index, ig_attr_edge,
                                      node_alpha=ig_attr_node)
  plt.show()

  # # Node and edge explainability
  # # ============================

  captum_model = to_captum(model, mask_type='node_and_edge',
                          output_idx=output_idx)

  ig = IntegratedGradients(captum_model)
  ig_attr_node, ig_attr_edge = ig.attribute(
      (data.x.unsqueeze(0), edge_mask.unsqueeze(0)), target=target,
      additional_forward_args=(data.edge_index), internal_batch_size=1)

  # Scale attributions to [0, 1]:
  ig_attr_node = ig_attr_node.squeeze(0).abs().sum(dim=1)
  ig_attr_node /= ig_attr_node.max()
  ig_attr_edge = ig_attr_edge.squeeze(0).abs()
  ig_attr_edge /= ig_attr_edge.max()

  # Visualize absolute values of attributions:
  ax, G = explainer.visualize_subgraph(output_idx, data.edge_index, ig_attr_edge,
                                      node_alpha=ig_attr_node)
  plt.show()

running it like this:

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(num_of_feats, output_dim)
      
    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        return x



num_of_nodes = 14
num_of_feats = 5
output_dim   = 10
device       = 'cpu'
adj          = np.random.randint(0,2,[num_of_nodes,num_of_nodes])
feats        = np.random.uniform(-1,1,[num_of_nodes, num_of_feats])

model      = GCN()
captum_explain_graph(model, feats, adj, 1, 'cpu')

2. Second solution

I also tried with GNNexplainer

from torch_geometric.nn import GNNExplainer


def gnn_explainer_func(model, node_idx, feats, adj):

  model                     = model.float()
  feats                     = torch.Tensor(feats)
  adj                       = torch.Tensor(adj)
  edge_index                = adj.nonzero(as_tuple=False).t()
  explainer                 = GNNExplainer(model, epochs=20)
  node_feat_mask, edge_mask = explainer.explain_node(node_idx, feats, edge_index)
  ax, G = explainer.visualize_subgraph(node_idx, edge_index, edge_mask, threshold=0.5)
  plt.show()

Issues:

1. Am I doing it correctly?
2. captum requires target in ig.attribute but I am using GCN as feature extractor so I don't have target values, I have only adj_matrix and feats per node, what should I use for target value in this case? If I use None (as per captum), it gives an error. What is the solution for this?
3. both ( captum and gnnexplainer ) are giving different outputs, which one is better for this type of problem statement?

Here is full working code:

colab link

[rusty1s]

1. I think this code looks good.
2. target is indeed a problem since it should only be set to None for scalar outputs, see here. Nonetheless, you might be able to utilize other explainability methods in captum that do not require targets (no captum expert though).
3. In usual, we found captum explanations to be better than the ones generated by GNNExplainer.

[monk1337]

@rusty1s Thanks for the response. I thought about one solution

model      = GCN()
log_logits = model(feats_x, edge_x)

torch.sum(log_logits, dim = 0) or torch.mean(log_logits, dim = 0)

it will convert the output to the scaler and then set target = None
does it make sense from a Graph perspective? ( if yes then sum or average?)
Is there any other hack to use captum without a target?

[rusty1s]

I think this doc explains it nicely, see here. In that case, I assume that torch.sum will yield better results than torch.mean (since it will preserve more information in the output).

[monk1337]

That was helpful, I have decided to go with torch.sum method however I am facing one issue, How to handle multiple output from graph model in to_captum function?

For example if my network looks like this:

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(num_of_feats, output_dim)
      
    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        return {'x'  : x, 
                'sum': torch.sum(x, 1)}

Since to_captum takes inputs and return single output, How to use this model with to_captum?

I found one solution, I tried but giving error. Please guide on this.

[rusty1s]

You could wrap your model into a model that combines your model and torch.sum. Would that work?

[monk1337]

Is it correct approach?

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(num_of_feats, output_dim)
      
    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        return x


class GCN_custom(torch.nn.Module):
  def __init__(self):
    super().__init__()
    self.model = GCN()
      
  def forward(self, x, edge_index):
    x = self.model(x, edge_index)
    return torch.sum(x, dim = 0)



gcn_sum = GCN_custom()
gcn_sum = gcn_sum.double()

captum_model = to_captum(gcn_sum, mask_type='edge', output_idx=1)
edge_mask    = torch.ones(data.num_edges, requires_grad=True, device=device)
ig           = IntegratedGradients(captum_model)
ig_attr_edge = ig.attribute(edge_mask.unsqueeze(0), target=None,
                              additional_forward_args=(data.x, data.edge_index),
                              internal_batch_size=1)

[rusty1s]

Yes, that‘s what I meant.