Reshaping pyg mini-batch

[Gaurangkarwande]

I am working on a problem where I construct a graph of images. Each node is an image and has dimensions 1x224x224. Hence if there are 36 nodes in a graph, a single graph will have dimensions 36x1x224x224. Now, I am working with a batch size of 100. Hence, the dimensions of input to the model are 100x36x1x224x224. However, pyg creates a single disconnected graph of all graphs in the batch. Hence the actual input to the model is of size 3600x1x224x224. Now I want to pass each of the nodes through a CNN encoder. Hence I reshape the the tensor inside the model to 100x36x1x224x224 and now pass a batch of nodes to the CNN encoder. The CNN encoder extracts a 1024 dimensional feature vector for each image/node. I then again reshape the CNN output as it was earlier, i.e. 3600x1024.

def forward(self,x): 
     x = x.view(batch_size, num_nodes, x.shape[1], x.shape[2], x.shape[3])   #100 x 36 x 1 x 224 x 224
     encoding = []
    for i in range(num_nodes):
         encoding[i] = self.regional_cnn(x[:,i,:,:,:])                                            # 100 x 1024
    encoding = torch.stack(encoding, dim=1)                                            # 3600 x 1024
    return encoding

I wanted to know whether the internally structure of my graph is maintained in all these reshapings. That is when reshape from 3600x1x224x224 to 100x36x1x224x224 , I am actually getting 36 images of size 1x224x224 in a single sample of the batch.

[rusty1s]

This looks correct to me, but I'm not entirely sure why you cannot simply process the CNN of shape [3600, 1, 224, 224]. Since you want to re-use the same CNN across both nodes and batch dimension, you can well stack them together and leverage better parallelization.

[Gaurangkarwande]

Thank you for confirming. I do not have enough compute resources to handle a batch size of 3600 within the CNN encoder.

[rusty1s]

But that is what you are currently doing anyways, right? You have two options here:

1. Decrease the batch_size to allow to train CNN and GNN jointly.
2. Detach the CNN from the computation graph, and compute image embeddings of nodes in a pre-processing step. This is commonly done in GNN literature by making use of a pre-trained CNN architecture.