Using linear layers on batched graphs

[dottipr]

Hi,

I am implementing a GNN model as an exercise to get familiar with Pytorch Geometric and I have a doubt about linear layers when using mini-batches. I've tried to find a solution online and on the other threads, but I couldn't find any.

I am using a customised data class for my graphs (which are basically bipartite graphs):

class P(Data):
    def __init__(self, edge_index, x_l, x_c, y):
        # edge_index is a bipartite adjacency matrix 
        # x_l is the feature vector of the l nodes
        # x_c is the feature vector of the c nodes
        # y is the label: 1 or 0

        super(P, self).__init__()
        # nodes features
        self.x_l = x_l
        self.x_c = x_c
        self.y = y
        
        self.num_l = x_l.size(0)
        self.num_c= x_c.size(0)
        
        # edges
        self.edge_index = edge_index
        self.adj_t = SparseTensor(row = edge_index[1],
                                  col = edge_index[0],
                                  sparse_sizes = [self.num_c, self.num_l]
                                 )
        ...
        
        self.num_nodes = self.num_l + self.num_c

    def __inc__(self, key, value):
        if key == 'edge_index':
            return torch.tensor([[self.x_l.size(0)], [self.x_c.size(0)]])
        else:
            return super().__inc__(key, value)

And I load them in mini-batches using dataloader:
loader = DataLoader(dataset, batch_size=batch_size, follow_batch=['x_l','x_c']).

In my model I have several linear layers that are applied on the feature vectors x_l and x_c (they are updated at different times) contained in my class MLP:

class MLP(nn.Module):
    def __init__(self, d_in = 128, d_outs = [128,128,128,128], transfer_fn = 'relu'):
        super(MLP, self).__init__()

        self.linears = nn.ModuleList([])
        for d_out in d_outs:
            self.linears.append(nn.Linear(d_in, d_out, bias=True))
            d_in = d_out

        if transfer_fn == 'relu':
            self.transfer = nn.ReLU()
        else:
            raise NotImplementedError

    def forward(self, x):        
        for linear in self.linears[:-1]:
            x = self.transfer(linear(x))

        return self.linears[-1](x)

Instances of the class MLP are then used in MessagaPassing layers where I've implemented only the message_and_aggregate method, since the network's functions are given as matrix multiplications that use the adjacency matrix adj_t.

However, if I am not mistaken, when using batches, the features matrices x_l and x_c get concatenated along the first dimension and for this reason when I forward them on a MLP, the output of a node feature can depend also on the features of nodes belonging to other graphs (unless I've misunderstood how linear layers works).

Is there a way to use linear layers in such a way that the output values of each graph's feature matrix depend only on its own initial value, instead of having outputs which mix values coming from different graphs? The graphs can have different numbers of nodes. Moreover, I am also using Pytorch's implementation of LSTM for some layers and I guess that the same problem occurs also in that case.

I am happy to provide more information if needed!

Cheers,
Prisca

[rusty1s]

This is not a problem. A linear layer will transform each node feature vector in isolation - the node feature dimension acts as the batch dimension here:

x = ...  # [num_nodes, 256]
lin = Linear(256, 512)
x = self.lin(x)  # [num_nodes, 512]

You don't have to worry about that (even for LSTMs), as long as the node dimension is used as the batch dimension.

[dottipr]

Thanks a lot for your help (and sorry for my late answer).

This is really relieving and should solve my problem.

Have a nice day,
Prisca