Explainer for graph classification

[shubhiambast]

Hi,

I am using TopKPooling for my graph classification task.

class GNNTopK(torch.nn.Module): 
    def __init__(self, in_dim, hidden_channels, out_dim, pooling_ratio):
        super().__init__()
        self.conv1 = GraphConv(in_dim, hidden_channels)
        self.pool1 = TopKPooling(hidden_channels, ratio=pooling_ratio)
        self.conv2 = GraphConv(hidden_channels, hidden_channels)
        self.pool2 = TopKPooling(hidden_channels, ratio=pooling_ratio)
        self.lin1 = torch.nn.Linear(hidden_channels*2, hidden_channels)
        self.lin2 = torch.nn.Linear(hidden_channels, hidden_channels//2)
        self.lin3 = torch.nn.Linear(hidden_channels//2, out_dim)
        
    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch
        
        # First layer
        x = F.relu(self.conv1(x, edge_index))
        x, edge_index, _, batch, perm1, _ = self.pool1(x, edge_index, None, batch)
        x1 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)
        
        # Second layer
        x = F.relu(self.conv2(x, edge_index))
        x, edge_index, _, batch, perm2, _ = self.pool2(x, edge_index, None, batch)
        x2 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)
        x = x1 + x2
        
        x = F.relu(self.lin1(x))
        x = F.dropout(x, p=0.4, training=self.training)
        x = F.relu(self.lin2(x))
        x = F.dropout(x, p=0.4, training=self.training)
        x = self.lin3(x)
        return x

Is it possible to use GNNExplainer on this kind of model? Any suggestions would be appreciated.