Error while passing list of aggregators to GINEConv

[shubhamjmeesho]

I an getting this error,

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-04065f95-870f-468a-b1c6-17dd03a71ed9/lib/python3.10/site-packages/torch_geometric/nn/conv/gin_conv.py:183, in GINEConv.forward(self, x, edge_index, edge_attr, size)
    181 x_r = x[1]
    182 if x_r is not None:
--> 183     out = out + (1 + self.eps) * x_r
    185 return self.nn(out)

RuntimeError: The size of tensor a (256) must match the size of tensor b (128) at non-singleton dimension 1

during forward pass of the following GINEConv layer -

GINEConv(
Sequential(Linear(128, 128),
            BatchNorm(128), ReLU(),
            Linear(128, 128), ReLU()),
aggr = ['mean', 'max'],
edge_dim = 6
)

I cross-checked, the input I am passing to this layer has a dimension of 128. Kindly tell what is wrong here.

[akihironitta]

Could you provide your complete code?

[shubhamjmeesho]

@akihironitta adding the full code -

class HeteroGNN(torch.nn.Module):
    def __init__(self, dim_in, dim_h, dim_out, num_layers, aggr1, edge_dim):
        super().__init__()

        metadata = (['actor1', 'actor2'], [('actor1', 'actor1_actor1_association', 'actor1'), ('actor2', 'actor2_actor2_association', 'actor2'), ('actor1', 'actor1_actor2_association', 'actor2'), ('actor2', 'rev_actor1_actor2_association', 'actor1')])


        self.lin_dict = torch.nn.ModuleDict()
        for node_type in metadata[0]:
            self.lin_dict[node_type] = Linear(dim_in[node_type], dim_h[0])

        self.convs = torch.nn.ModuleList()
        for idx in range(num_layers):
            in_nodes = dim_h[0]*((len(metadata[0])*len(aggr1))**(idx))
            out_nodes = in_nodes*len(metadata[0])*len(aggr1)
            conv = HeteroConv({
                ('actor1', 'actor1_actor1_association', 'actor1'): GINEConv(
                                                            Sequential(Linear(in_nodes, in_nodes),
                                                                        BatchNorm(in_nodes), ReLU(),
                                                                        Linear(in_nodes, in_nodes), ReLU()),
                                                            aggr = aggr1,
                                                            edge_dim = edge_dim[0]
                                                            ),
                ('actor2', 'actor2_actor2_association', 'actor2'): GINEConv(
                                                            Sequential(Linear(in_nodes, in_nodes),
                                                                        BatchNorm(in_nodes), ReLU(),
                                                                        Linear(in_nodes, in_nodes), ReLU()),
                                                            aggr = aggr1,
                                                            edge_dim = edge_dim[1]
                                                            ),
                ('actor1', 'actor1_actor2_association', 'actor2'): GINEConv(
                                                            Sequential(Linear(in_nodes, in_nodes),
                                                                        BatchNorm(in_nodes), ReLU(),
                                                                        Linear(in_nodes, in_nodes), ReLU()),
                                                            aggr = aggr1,
                                                            edge_dim = edge_dim[2]
                                                            ),
                ('actor2', 'rev_actor1_actor2_association', 'actor1'): GINEConv(
                                                            Sequential(Linear(in_nodes, in_nodes),
                                                                        BatchNorm(in_nodes), ReLU(),
                                                                        Linear(in_nodes, in_nodes), ReLU()),
                                                            aggr = aggr1,
                                                            edge_dim = edge_dim[3]
                                                            )                
            }, aggr="cat")
            self.convs.append(conv)

        self.lin1 = Linear(out_nodes, dim_h[1])
        self.lin2 = Linear(dim_h[1], dim_out)

        self.optimizer = torch.optim.Adam(self.parameters(),
                                        lr=0.01)

    def forward(self, x_dict, edge_index_dict, edge_attr_dict):

        for node_type, x in x_dict.items():
            x_dict[node_type] = self.lin_dict[node_type](x).relu()

        out = [x_dict]
        for conv in self.convs:
            out_next = conv(out[-1], edge_index_dict, edge_attr_dict)
            out.append(out_next)

        h = out[-1]['actor1'] 

        l1 = self.lin1(h).relu()
        drop_l1 = F.dropout(l1, p=0.2, training=self.training)
        l2 = self.lin2(drop_l1)

        return F.log_softmax(l2, dim=1), l1


graph = torch.load('path to HeteroData object')

train_loaders = NeighborLoader(graph,
                                    num_neighbors=[10, 5],  
                                    batch_size=2048, 
                                    shuffle = True,
                                    drop_last = True,
                                    input_nodes=('actor1', graph['actor1'].train_mask),
                                    )

edge_dim = []
edge_dim.append(graph.num_edge_features[('actor1', 'actor1_actor1_association', 'actor1')])
edge_dim.append(graph.num_edge_features[('actor2', 'actor2_actor2_association', 'actor2')])
edge_dim.append(graph.num_edge_features[('actor1', 'actor1_actor2_association', 'actor2')])
edge_dim.append(graph.num_edge_features[('actor2', 'rev_actor1_actor2_association', 'actor1')])
heterognn = HeteroGNN(graph.num_features, [128,32], 2, 1, ["mean", "max"], edge_dim)
heterognn = heterognn.to(device)

criterion = torch.nn.CrossEntropyLoss()
optimizer = graph.optimizer
heterognn.train()
for epoch in range(100):
    # initiate tracking metrics
    total_loss = 0
    total_train_examples = 0


    # Train on batches from all graphs
    for batch in train_loader:
        optimizer.zero_grad()
        batch = batch.to(device)
        batch_size = batch['actor1'].batch_size
        total_train_examples += batch_size

        # Forward pass
        soft_l2, l1 = heterognn(batch.x_dict, batch.edge_index_dict, batch.edge_attr_dict)

        # compute loss and accuracy
        loss = criterion(soft_l2[batch['actor1'].train_mask], batch['actor1'].y[batch['actor1'].train_mask]) 
        total_loss += loss* batch_size

        # Back propagation
        loss.backward()
        optimizer.step()

[rusty1s]

I haven't run the example but it looks like in_nodes is multiplied by the number of node types, although the convolutions in HeteroConv are operating on a single node type only.

[shubhamjmeesho]

I have done this because I have used cat as the aggregator for heteroconv layer. Although this would not affect anything as I am using single conv layer only

[rusty1s]

The input to the first HeteroConv module would still expect non-concatenated inputs.

[shubhamjmeesho]

It has to the power idx and idx will be zero in first layer

[rusty1s]

Ah, missed this. I will try to reproduce.