How to explain the GNN network of MessagePassing class

[688spiderbin]

Hello!
Could you update the tutorial of Captum for network of MessagePassing class in the doc. I have tried searching some discussions about how to explain the messagepassing class. But it's quite messy and it's hard to follow.

I appreciate a lot !

[rusty1s]

Can you clarify what you want to have updated exactly, and where you are stuck?

[rusty1s]

How is your CaptumExplainer defined?

[688spiderbin]

I remember I upload that screen shot before, please check :

[rusty1s]

Note that you need to pass in additional arguments as keyword arguments in the explainer:

explainer(data.x, data.edge_index, edge_attr=data.edge_attr, batch=data.batch)

I am not sure if this resolved your issue.

[688spiderbin]

Yep, you have suggested to me this before, it's cannot work.

[rusty1s]

Sorry for late reply. This works fine on our end, see here. I am not yet sure what is the problem though.

[688spiderbin]

Thanks for your recommendation, I rechecked the model and can't found any parameter haven't use. Indeed It's hard to detect the issue and describe it exactly, so I paste all the code here.

import torch
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree, softmax
from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool, GlobalAttention, Set2Set
import torch.nn.functional as F
from torch_scatter import scatter_add
from torch_geometric.nn.inits import glorot, zeros

The model is

class GINConv(MessagePassing):
    def __init__(self, emb_dim, aggr = "add"):
        super(GINConv, self).__init__()
        #multi-layer perceptron
        self.mlp = torch.nn.Sequential(torch.nn.Linear(emb_dim, 2*emb_dim), torch.nn.ReLU(), torch.nn.Linear(2*emb_dim, emb_dim))
        self.edge_embedding1 = torch.nn.Embedding(num_bond_type, emb_dim)
        self.edge_embedding2 = torch.nn.Embedding(num_bond_direction, emb_dim)

        torch.nn.init.xavier_uniform_(self.edge_embedding1.weight.data)
        torch.nn.init.xavier_uniform_(self.edge_embedding2.weight.data)
        self.aggr = aggr

    def forward(self, x, edge_index, edge_attr):
        #add self loops in the edge space
        edge_index = add_self_loops(edge_index, num_nodes = x.size(0))[0]

        #add features corresponding to self-loop edges.
        self_loop_attr = torch.zeros(x.size(0), 2)
        self_loop_attr[:,0] = 4 #bond type for self-loop edge
        self_loop_attr = self_loop_attr.to(edge_attr.device).to(edge_attr.dtype)
        edge_attr = torch.cat((edge_attr, self_loop_attr), dim = 0)

        edge_embeddings = self.edge_embedding1(edge_attr[:,0]) + self.edge_embedding2(edge_attr[:,1])

        # return self.propagate(self.aggr, edge_index, x=x, edge_attr=edge_embeddings)
        return self.propagate(edge_index, x=x, edge_attr=edge_embeddings)

    def message(self, x_j, edge_attr):
        return x_j + edge_attr

    def update(self, aggr_out):
        return self.mlp(aggr_out)

class GNN(torch.nn.Module):
   """
    Args:
        num_layer (int): the number of GNN layers
        emb_dim (int): dimensionality of embeddings
        JK (str): last, concat, max or sum.
        max_pool_layer (int): the layer from which we use max pool rather than add pool for neighbor aggregation
        drop_ratio (float): dropout rate
        gnn_type: gin, gcn, graphsage, gat

    Output:
        node representations

    """
    def __init__(self, num_layer, emb_dim, JK = "last", drop_ratio = 0, gnn_type = "gin"):
        super(GNN, self).__init__()
        self.num_layer = num_layer
        self.drop_ratio = drop_ratio
        self.JK = JK

        if self.num_layer < 2:
            raise ValueError("Number of GNN layers must be greater than 1.")

        self.x_embedding1 = torch.nn.Embedding(num_atom_type, emb_dim)
        self.x_embedding2 = torch.nn.Embedding(num_chirality_tag, emb_dim)

        torch.nn.init.xavier_uniform_(self.x_embedding1.weight.data)
        torch.nn.init.xavier_uniform_(self.x_embedding2.weight.data)

        ###List of MLPs
        self.gnns = torch.nn.ModuleList()
        for layer in range(num_layer):
            if gnn_type == "gin":
                self.gnns.append(GINConv(emb_dim, aggr = "add"))
            elif gnn_type == "gcn":
                self.gnns.append(GCNConv(emb_dim))
            elif gnn_type == "gat":
                self.gnns.append(GATConv(emb_dim))
            elif gnn_type == "graphsage":
                self.gnns.append(GraphSAGEConv(emb_dim))


        ###List of batchnorms
        self.batch_norms = torch.nn.ModuleList()
        for layer in range(num_layer):
            self.batch_norms.append(torch.nn.BatchNorm1d(emb_dim))

    def forward(self, x, edge_index, edge_attr):
        x = self.x_embedding1(x[:, 0].to(dtype=torch.int64)) + self.x_embedding2(x[:, 1].to(dtype=torch.int64))
        h_list = [x]
        for layer in range(self.num_layer):
            h = self.gnns[layer](h_list[layer], edge_index, edge_attr)
            h = self.batch_norms[layer](h)
            #h = F.dropout(F.relu(h), self.drop_ratio, training = self.training)
            if layer == self.num_layer - 1:
                #remove relu for the last layer
                h = F.dropout(h, self.drop_ratio, training = self.training)
            else:
                h = F.dropout(F.relu(h), self.drop_ratio, training = self.training)
            h_list.append(h)

        ### Different implementations of Jk-concat
        if self.JK == "concat":
            node_representation = torch.cat(h_list, dim = 1)
        elif self.JK == "last":
            node_representation = h_list[-1]
        elif self.JK == "max":
            h_list = [h.unsqueeze_(0) for h in h_list]
            node_representation = torch.max(torch.cat(h_list, dim = 0), dim = 0)[0]
        elif self.JK == "sum":
            h_list = [h.unsqueeze_(0) for h in h_list]
            node_representation = torch.sum(torch.cat(h_list, dim = 0), dim = 0)[0]

        return node_representation

class GNN_graphpred(torch.nn.Module):
    """
    Extension of GIN to incorporate edge information by concatenation.

    Args:
        num_layer (int): the number of GNN layers
        emb_dim (int): dimensionality of embeddings
        num_tasks (int): number of tasks in multi-task learning scenario
        drop_ratio (float): dropout rate
        JK (str): last, concat, max or sum.
        graph_pooling (str): sum, mean, max, attention, set2set
        gnn_type: gin, gcn, graphsage, gat
        
    See https://arxiv.org/abs/1810.00826
    JK-net: https://arxiv.org/abs/1806.03536
    """
    def __init__(self, num_layer, emb_dim, num_tasks, gnn_type, JK = "last", drop_ratio = 0, graph_pooling = "mean"):
        super(GNN_graphpred, self).__init__()
        self.num_layer = num_layer
        self.drop_ratio = drop_ratio
        self.JK = JK
        self.emb_dim = emb_dim
        self.num_tasks = num_tasks

        if self.num_layer < 2:
            raise ValueError("Number of GNN layers must be greater than 1.")

        self.gnn = GNN(num_layer, emb_dim, JK, drop_ratio, gnn_type = gnn_type)

        #Different kind of graph pooling
        if graph_pooling == "sum":
            self.pool = global_add_pool
        elif graph_pooling == "mean":
            self.pool = global_mean_pool
        elif graph_pooling == "max":
            self.pool = global_max_pool
        elif graph_pooling == "attention":
            if self.JK == "concat":
                self.pool = GlobalAttention(gate_nn = torch.nn.Linear((self.num_layer + 1) * emb_dim, 1))
            else:
                self.pool = GlobalAttention(gate_nn = torch.nn.Linear(emb_dim, 1))
        elif graph_pooling[:-1] == "set2set":
            set2set_iter = int(graph_pooling[-1])
            if self.JK == "concat":
                self.pool = Set2Set((self.num_layer + 1) * emb_dim, set2set_iter)
            else:
                self.pool = Set2Set(emb_dim, set2set_iter)
        else:
            raise ValueError("Invalid graph pooling type.")

        #For graph-level binary classification
        if graph_pooling[:-1] == "set2set":
            self.mult = 2
        else:
            self.mult = 1
        
        if self.JK == "concat":
            self.graph_pred_linear = torch.nn.Linear(self.mult * (self.num_layer + 1) * self.emb_dim, self.num_tasks)
        else:
            self.graph_pred_linear = torch.nn.Linear(self.mult * self.emb_dim, self.num_tasks)

    def from_pretrained(self, model_file):
        #self.gnn = GNN(self.num_layer, self.emb_dim, JK = self.JK, drop_ratio = self.drop_ratio)
        self.gnn.load_state_dict(torch.load(model_file, map_location='cpu'))

    def forward(self, x, edge_index, edge_attr, batch):
        node_representation = self.gnn(x, edge_index, edge_attr)
        return self.graph_pred_linear(self.pool(node_representation, batch)) 

And the train function is

criterion = nn.BCEWithLogitsLoss(reduction = "none")

def train(args, model, device, loader, optimizer):
    model.train()

    for step, batch in enumerate(tqdm(loader, desc="Iteration")):
        batch = batch.to(device)
        pred = model(batch.x, batch.edge_index, batch.edge_attr, batch.batch)
        y = batch.y.view(pred.shape).to(torch.float64)

        #Whether y is non-null or not.
        is_valid = y**2 > 0
        #Loss matrix
        loss_mat = criterion(pred.double(), (y+1)/2)
        #loss matrix after removing null target
        loss_mat = torch.where(is_valid, loss_mat, torch.zeros(loss_mat.shape).to(loss_mat.device).to(loss_mat.dtype))
            
        optimizer.zero_grad()
        loss = torch.sum(loss_mat)/torch.sum(is_valid)
        loss.backward()
        optimizer.step()

The code for explanation

# set-up dataset
dataset = MoleculeDataset("E:/github_download/SourceCode/pretrain-gnns-master/dataset/" + 'mutag', dataset='mutag')

train_dataset, valid_dataset = cv_random_split(dataset, seed=42)
print(train_dataset[0])
data = ([t for t in train_dataset])[0]  # choose one 
print(data)


# using pre-trained model
input_model_file = "E:/github_download/SourceCode/pretrain-gnns-master/chem/model_gin/supervised_masking.pth"
model = GNN_graphpred(5, 300, 1, gnn_type='gin', JK='last',
                      drop_ratio=0.5,
                      graph_pooling="mean")
model.from_pretrained(input_model_file)


model.to(device)
model.eval()
data = data.to(device)
pred = model(data.x, data.edge_index, data.edge_attr, data.batch)
print(pred)


explainer = Explainer(
    model = model,
    algorithm=CaptumExplainer('IntegratedGradients'),  
    explanation_type='model',
    model_config=dict(
        mode = 'binary_classification',
        task_level = 'graph',
        return_type = 'probs',
    ),
    node_mask_type='attributes',
    edge_mask_type='object',
    threshold_config=dict(
        threshold_type='topk',
        value=200,
    ),
)


explanation = explainer(data.x, data.edge_index, edge_attr=data.edge_attr, batch=data.batch)
print(explanation.available_explanations)

path = 'feature_importance.png'
explanation.visualize_feature_importance(path, top_k=10)
print(f"Feature importance plot has been saved to '{path}'")

path = 'subgraph.pdf'
explanation.visualize_graph(path)
print(f"Subgraph plot has been saved to '{path}'")

Then, it comes up

Traceback (most recent call last):
  File "E:\github_download\SourceCode\pretrain-gnns-master\chem\explain(成功).py", line 64, in <module>
    explanation = explainer(data.x, data.edge_index, edge_attr=data.edge_attr, batch=data.batch)
  File "E:\anaconda3\envs\py38\lib\site-packages\torch_geometric\explain\explainer.py", line 205, in __call__
    explanation = self.algorithm(
  File "E:\anaconda3\envs\py38\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "E:\anaconda3\envs\py38\lib\site-packages\torch_geometric\explain\algorithm\captum_explainer.py", line 170, in forward
    attributions = self.attribution_method_instance.attribute(
  File "E:\anaconda3\envs\py38\lib\site-packages\captum\log\__init__.py", line 42, in wrapper
    return func(*args, **kwargs)
  File "E:\anaconda3\envs\py38\lib\site-packages\captum\attr\_core\integrated_gradients.py", line 274, in attribute
    attributions = _batch_attribution(
  File "E:\anaconda3\envs\py38\lib\site-packages\captum\attr\_utils\batching.py", line 78, in _batch_attribution
    current_attr = attr_method._attribute(
  File "E:\anaconda3\envs\py38\lib\site-packages\captum\attr\_core\integrated_gradients.py", line 351, in _attribute
    grads = self.gradient_func(           # 在这里计算梯度
  File "E:\anaconda3\envs\py38\lib\site-packages\captum\_utils\gradient.py", line 119, in compute_gradients
    grads = torch.autograd.grad(torch.unbind(outputs), inputs)
  File "E:\anaconda3\envs\py38\lib\site-packages\torch\autograd\__init__.py", line 303, in grad
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.

[688spiderbin]

AND Could you help me to find out which parameters haven't used according to the code? Is it torch.nn.Embedding ? I don't think this will be the answer.

>>>self.x_embedding1(x[:, 0].to(dtype=torch.int64)) + self.x_embedding2(x[:, 1].to(dtype=torch.int64))
Out[1]: 
tensor([[ 1.4876, -0.7566,  3.8224,  ..., -3.0525, -2.4919, -1.2139],
        [ 1.4876, -0.7566,  3.8224,  ..., -3.0525, -2.4919, -1.2139],
        [ 1.4876, -0.7566,  3.8224,  ..., -3.0525, -2.4919, -1.2139],
        ...,
        [ 0.1232, -0.7887,  3.5018,  ...,  3.5915,  1.4795, -1.6173],
        [-2.2018,  0.8337,  3.6481,  ...,  2.3498, -0.4928,  0.4342],
        [-2.2018,  0.8337,  3.6481,  ...,  2.3498, -0.4928,  0.4342]],
       grad_fn=<AddBackward0>)