Custom Conv

[DogitoErgoSum]

Hello,

I am building an Heterogeneous graph classification model, all my graphs have at max 1 node per node type. This means i can improve a lot the gineconv expressiveness by doing a convolution that does the following:
MLP(concat( i,j)) ,where i is the incident node and j a neighbor

I am trying to create a convolution class that does that like this:

import torch
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree
from torch_geometric.nn import MLP
class GCNConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr='add')  
        
        mlp = MLP([-1, 32, 16,1])
    def forward(self, x, edge_index):
        
        return self.propagate(edge_index, x=x)

    def message(self, x_i, x_j):
        
        return mlp(torch.concat([x_i,x_j],axis=1))

Am i doing something wrong?

[Padarn]

Can you elaborate on what you think might be wrong?

Though there ae some bugs in your code - for example you will want to set the mlp as a property of your network

self.mlp = MLP([-1, 32, 16,1])

[DogitoErgoSum]

That helped! It just seemed to to simple to be true, but now is kinda working.
I have the following bug with MLP()
ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 32])
i m working with batches of 512 HeteroData graphs, but some nodes types will have only 1 node per batch.
Is there a solution to this problem? Or do i need to build my own neural net in def init ?

[DogitoErgoSum]

I can build the neural net but i m not sure if it works with the batch

[rusty1s]

The only solution I can think of is to disable batch normalization inside the MLP. It is not well-defined for only normalizing a single "node".