Extracting the individual messages for a message passing network

[CagkanYapar]

Hi,

I have a question regarding message passing networks. I am interested in getting the individual messages, i.e. the edge values and not only the node values.
Is there a way to get the individual messages for each edge out? Is it possible to do this by a smart modification of message and update functions?

[rusty1s]

There exists multiple options to do this:

• You can use hooks, see here.
• You can first compute the message via MessagePassing.edge_updater, and input the computed value into MessagePassing.propagate (basically disentangling node and value computations)
• You can store intermediate results as a attribute

def forward(self, ...):
    self.propagate(edge_index, ...)
    edge_attr = self.my_edge_attr
    self.my_edge_attr = None
    
def message(self, ...):
    self.my_edge_attr = ...

[CagkanYapar]

Thank you for the answer!

Is to be computed message first declared by MessagePassing.edge_update ?

Could you show this on a specific example, e.g. edge convolutional layer? How could one modify it such that the messages would be output? How can one then access them, via the edge_index? The code for edge convolution on https://pytorch-geometric.readthedocs.io/en/latest/notes/create_gnn.html reads:

`import torch
from torch.nn import Sequential as Seq, Linear, ReLU
from torch_geometric.nn import MessagePassing

class EdgeConv(MessagePassing):
def init(self, in_channels, out_channels):
super().init(aggr='max') # "Max" aggregation.
self.mlp = Seq(Linear(2 * in_channels, out_channels),
ReLU(),
Linear(out_channels, out_channels))

def forward(self, x, edge_index):
    # x has shape [N, in_channels]
    # edge_index has shape [2, E]

    return self.propagate(edge_index, x=x)

def message(self, x_i, x_j):
    # x_i has shape [E, in_channels]
    # x_j has shape [E, in_channels]

    tmp = torch.cat([x_i, x_j - x_i], dim=1)  # tmp has shape [E, 2 * in_channels]
    return self.mlp(tmp)`

[rusty1s]

Here is a simple example: https://github.com/pyg-team/pytorch_geometric/blob/master/test/nn/conv/test_message_passing.py#L228