NodeModel for undirected edges in MetaLayer

[arunppsg]

Hi all,

Just for clarification, the NodeModel in MetaLayer docs is for directed graphs right?

For using the NodeModel in undirected graph, I wrote the following snippet:

class NodeModel(nn.Module):
    
    def __init__(self, nfeat_edge=32, nfeat_node=32, nfeat_global=32):
        super(NodeModel, self).__init__()
        # A single layer for node processing
        self.node_mlp = nn.Sequential(
                nn.Linear(in_features=nfeat_edge + nfeat_node + nfeat_global,
                         out_features=nfeat_node)
            )
        
    def forward(self, x, edge_index, edge_attr, global_state):
        # edge_index represents the connectivity in the graph
        # Shape: (2, num_edges) i.e for each edge, it holds start and end vertex
        # Computing mean edge attributes
        row, col = edge_index
        a = edge_attrs[row]
        b = edge_attrs[col]
        mean_edge_feats = scatter_mean(torch.cat((a, b)), torch.cat((row, col)), dim=0)
        out = torch.cat([x, mean_edge_feats, global_state], dim=1)
        return self.node_mlp(out)

This will later be passed into MetaLayer initialization. Am I right in assuming out = torch.cat([x, mean_edge_feats, global_state], dim=1)?

[rusty1s]

In general, for undirected graphs, it is best to simply hold bidirectional edges in edge_index. Your approach will still account for the direction of edges (given that a and b might be different).

[arunppsg]

Thanks for your pointers. I am redefining the forward func. as follows for node update:

def forward(self, x, edge_index, edge_attr, global_state):
    # Holding bidirectional edges in edge_index and update edge_attrs
    edge_index = torch.cat((edge_index, torch.index_select(edge_index, 0, torch.LongTensor([1, 0]))), axis=1)
    row, col = edge_index
    edge_attrs = torch.cat((edge_attrs, edge_attrs), axis=0)
    # Computing node wise edge means
    out = scatter_mean(edge_attrs, col, dim=0)
    out = torch.cat([x, out, u[batch]], dim=1)
    return self.node_mlp(out)

Kindly let me know in case I am wrong here. Thanks!

[rusty1s]

This looks good. I think you are missing the aggregation of neighboring node features though:

def forward(self, x, edge_index, edge_attr, global_state):
    # Holding bidirectional edges in edge_index and update edge_attrs
    edge_index = torch.cat((edge_index, edge_index.flip([0])), axis=1)
    row, col = edge_index
    edge_attr = torch.cat((edge_attr, edge_attr), axis=0)
    out = torch.cat([x[row], edge_attr], dim=1)
    out = self.node_mlp_1(out)
    out = scatter_mean(out, col, dim=0, dim_size=x.size(0))
    out = torch.cat([x, out, u[batch]], dim=1)
    return self.node_mlp_2(out)

[arunppsg]

Yeah, thanks for your help!