Confusing with PyG pre-defined models and data shapes.

[xdhcode]

Dataset introduction

I'm trying to train pre-defined models to predict my dataset.
One sample is like Data(x=[48, 1], edge_index=[2, 53], y=0), and there are 615 samples constructing a DataLoader like DataBatch(x=[29520, 1], edge_index=[2, 32595], y=[615], batch=[29520], ptr=[616]). Data list is used to construct the DataLoader as below.

train_list=[]
for i in range(len(y_train)):
    X=torch.tensor(x_train[:,i]).float()
    Y=torch.tensor(y_train[i,0]).long()
    E=torch.tensor(e).long()
    data = Data(x=X, edge_index=E,y=Y)
    train_list.append(data)
train_loader= DataLoader(train_list, batch_size=len(y_train), shuffle=False)
#len(y_train): 615

Pre-defined model I used

I tried models.GCN(in_channels=data.num_node_features #1, out_channels=len(np.unique(y_train)) #53 ), here 53 labels for my dataset.

Problem encountered

I got the input shape torch.Size([29520, 1]) and output shape torch.Size([29520, 53]). And loss function nn.CrossEntropyLoss() can not deal with this shape, reporting Expected input batch_size (29520) to match target batch_size (615) . The correct output shape is supposed to be torch.Size(615, 53]).

User-defined model can work

In my self-defined model, the tensor shape is reshaped like x=x.reshape([int(len(x)/48), 48*hidden_channels]) before output linear layer, so that input shape torch.Size([29520, 1]) is transformed to torch.Size([615, 48*hidden_channels]). The output layer shape is [48*hidden_channels, 53], and that's how I make it work.

Question

How to make pre-defined model output the shape that the loss function can accept? Or, are other parts wrong? I can't figure out because the old-fashion way works. I really appreciate your advice.

[rusty1s]

It looks like you are missing a global readout step. By default, the basic GNNs we provide will return node-level representations. You can fix this, e.g., by running

class MyGNN(torch.nn.Module):
    def __init__(self, ...)
         self.gnn = GCN(...)
         self.classifier = Lin(...)
         
     def forward(self, x, edge_index, batch):
          x = self.gnn(x, edge_index)
          x = global_mean_pool(x, batch)
          return self.classifier(x)

[xdhcode]

Thanks, it works. And I notice that models.GCN provides out_channels to apply a final linear transformation to convert hidden node embeddings to output size. Is this supposed to work as a classifier and transform data shape in the manner introduced in User-defined model can work? Why do I still need global_mean_pool()?

[rusty1s]

The out_channels would only transforms the node embeddings, e.g., for node-level tasks. It is not related to global readouts/pooling.

[xdhcode]

I tried models.MLP to compare with PyTorch MLP I built, but it provides in_channels instead of in_features. To my knowledge, if nodes have one feature, then in_channels=1 for models.MLP, and in_features=num_nodes for nn.Linear(). Is this mean that models.MLP provides multiple layers like convolution kernels? Then how to define the hidden unit nums? So confusing. What's wrong with my attempt?

[rusty1s]

in_channels and in_features are meant to be the same.