Explainer for MetaLayer

[lfangyu09]

Hi,

Thank you so much for your efforts in this library. I want to use Explainer for MetaLayer (node regression), but I faced an issue. The output of MetaLayer is a tuple. There is an error AttributeError: 'tuple' object has no attribute 'size'. I tried explanation_type='model' and explanation_type='phenomenon', but all failed. Therefore, may I ask how to use Explainer for MetaLayer (node regression)? Thanks a lot.

[wsad1]

You'll have to write a small wrapper around the meta layer to not return a tuple of x, edge_attr and u, but just x.

class Wrapper(MetaLayer)
    def forward(...) same arguments as metalayer
         x, _, _ = super().forward(...)
         return x

[lfangyu09]

Thanks a lot. I tried but it gave me another issue.

class GNN_sensitivity(Module):
  def __init__(self, model_path=None):
    super(GNN_sensitivity, self).__init__()
    self.model_path = model_path
  def forward(self, x, edge_index, edge_attr):
    simulator = GNN()
    simulator = simulator.cuda()
    checkpoint = torch.load(self.model_path)
    simulator.load_state_dict(checkpoint["model"])
    out, _ = simulator.graph_processor(x, edge_index, edge_attr)
    return out

GNN_sensitivity_analysis = GNN_sensitivity(model_path)
data_graph_x_check = GNN_sensitivity_analysis(data_x, data.edge_index, data_edge_attr)

[rusty1s]

Can you share a small reproducible example? In the end, I don't think the MetaLayer is applicable in the Explainer module, since it doesn't make use of MessagePassing modules.

[lfangyu09]

Hi,

I used the code from the website

from torch import nn, cat
from torch_geometric.nn import MetaLayer
from torch_scatter import scatter_sum

class MLP(nn.Module):
    #MLP with LayerNorm
    def __init__(self, 
            in_dim, 
            out_dim=128, 
            hidden_dim=128,
            hidden_layers=2,
            norm_type='LayerNorm'):
        super(MLP, self).__init__()

        layers = [nn.Linear(in_dim, hidden_dim), nn.ReLU()]
        for _ in range(hidden_layers - 1):
            layers += [nn.Linear(hidden_dim, hidden_dim), nn.ReLU()]
        layers.append(nn.Linear(hidden_dim, out_dim))

        if norm_type is not None:
            assert (norm_type in ['LayerNorm', 'GraphNorm', 'InstanceNorm', 'BatchNorm', 'MessageNorm'])
            norm_layer = getattr(nn, norm_type)
            layers.append(norm_layer(out_dim))

        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)

#############################
class EdgeProcessor(nn.Module):

    def __init__(self, 
            in_dim_node=128, in_dim_edge=128,
            hidden_dim=128, 
            hidden_layers=2,
            norm_type='LayerNorm'):
        super(EdgeProcessor, self).__init__()
        self.edge_mlp = MLP(2 * in_dim_node + in_dim_edge, 
            in_dim_edge, 
            hidden_dim,
            hidden_layers,
            norm_type)

    def forward(self, src, dest, edge_attr, u=None, batch=None):
        out = cat([src, dest, edge_attr], -1) #concatenate source node, destination node, and edge embeddings
        out = self.edge_mlp(out)
        out += edge_attr #residual connection

        return out

class NodeProcessor(nn.Module):
    def __init__(self, 
            in_dim_node=128, in_dim_edge=128,
            hidden_dim=128, 
            hidden_layers=2,
            norm_type='LayerNorm'):
        super(NodeProcessor, self).__init__()
        self.node_mlp = MLP(in_dim_node + in_dim_edge,  
            in_dim_node,
            hidden_dim,
            hidden_layers,
            norm_type)

    def forward(self, x, edge_index, edge_attr, u=None, batch=None):
        row, col = edge_index
        out = scatter_sum(edge_attr, col, dim=0) #aggregate edge message by target
        out = cat([x, out], dim=-1)
        out = self.node_mlp(out)
        out += x #residual connection

        return out

def build_graph_processor_block(in_dim_node=128, in_dim_edge=128,
        hidden_dim_node=128, hidden_dim_edge=128, 
        hidden_layers_node=2, hidden_layers_edge=2,
        norm_type='LayerNorm'):
    return MetaLayer(
            edge_model=EdgeProcessor(in_dim_node, in_dim_edge, hidden_dim_edge, hidden_layers_edge, norm_type),
            node_model=NodeProcessor(in_dim_node, in_dim_edge, hidden_dim_node, hidden_layers_node, norm_type)
        )

class GraphProcessor(nn.Module):
    def __init__(self, 
        mp_iterations=15, 
        in_dim_node=128, in_dim_edge=128,
        hidden_dim_node=128, hidden_dim_edge=128, 
        hidden_layers_node=2, hidden_layers_edge=2,
        norm_type='LayerNorm'):
        super(GraphProcessor, self).__init__()

        self.blocks = nn.ModuleList()
        for _ in range(mp_iterations):
            self.blocks.append(build_graph_processor_block(in_dim_node, in_dim_edge,
                hidden_dim_node, hidden_dim_edge, 
                hidden_layers_node, hidden_layers_edge, 
                norm_type))

    def forward(self, x, edge_index, edge_attr):
        for block in self.blocks:
            x, edge_attr, _ = block(x, edge_index, edge_attr)

        return x, edge_attr

simulator = GraphProcessor()    
class GNN_sensitivity(nn.Module):
    def __init__(self, model):
        super(GNN_sensitivity, self).__init__()
        self.model = model
    def forward(self, x, edge_index, edge_attr):
        out, _ = self.model(x, edge_index, edge_attr)
        return out

GNN_sensitivity_analysis = GNN_sensitivity(model=simulator)
data_x = torch.rand(4, 128)
data_edge_index = torch.tensor([[0,1,2,3,3,2,3,1],
                  [3,2,3,1,0,1,2,3,]])
data_edge_attr = torch.rand(8, 128)
data_graph_x_check = GNN_sensitivity_analysis(data_x, data_edge_index, data_edge_attr)

[rusty1s]

Writing NodeProcessor as

class NodeProcessor(MessagePassing):
    def __init__(self, in_dim_node=128, in_dim_edge=128, hidden_dim=128,
                 hidden_layers=2, norm_type='LayerNorm'):
        super(NodeProcessor, self).__init__(aggr='sum')
        self.node_mlp = MLP(in_dim_node + in_dim_edge, in_dim_node, hidden_dim,
                            hidden_layers, norm_type)

    def forward(self, x, edge_index, edge_attr, u=None, batch=None):
        out = self.propagate(edge_index, edge_attr=edge_attr)

        out = cat([x, out], dim=-1)
        out = self.node_mlp(out)
        out += x  # residual connection

        return out

    def message(self, edge_attr):
        return edge_attr

and calling via

explainer = Explainer(
    model=GNN_sensitivity_analysis,
    algorithm=GNNExplainer(epochs=200),
    explanation_type='model',
    node_mask_type='attributes',
    edge_mask_type='object',
    model_config=dict(
        mode='regression',
        task_level='node',
        return_type='raw',
    ),
)
explainer(data_x, data_edge_index, edge_attr=data_edge_attr)

works for me.

[lfangyu09]

Thank you so much.