How to make batch with pictures of different sizes for model in PyTorch?

Viewed 433

I want to use GlobalAveragePooling in my PyTorch model and not to resize, crop or pad the image. I can train my model using only one image every iteration (not batch). But it is too slow and I don't know how to use several images of different sizes as one input for Model. Example of model code:

class GAPModel(nn.Module):
  def __init__(self):
    super().__init__()
    
    self.conv = nn.Sequential(
        nn.Conv2d(3, 16, kernel_size=3),
        nn.ReLU(inplace=True),
    )
    
    self.linear = nn.Sequential(
        nn.Linear(in_features=16, out_features=1),
        nn.ReLU(),
    )

  def forward(self, image):
    return self.linear(self.conv(image).mean([2, 3]))
2 Answers

As you already know, when you train a model you provide it with your input data.
The shape of your input data looks like this: channel_last: (batch_size, height, width, num_of_channels) or channel_first: (batch_size, num_of_channels, height, width)
If batch_size=32 (input_data.shape = (32,height,width,channels) , it means that all 32 images should have the same height, width, and num_of_channels.
Because you cannot store your input_data with different shapes into one numpy_array or tensor.
More clarification:
For example, if you have two different images, img_1 and img_2.
Their shapes are image1.shape=224,224,3, image2.shape=288,288,3, and you want to make batch_of_imgs with these two images. If so, what will be the shape of batch_of_imgs ?
Will it be batch_of_imgs.shape=(2,224,224,3) or (2,288,288,3) ?
As you see, when making a batch of images they should be the same shape.
It means that your input data(numpy_arrays or tensors) should be rectangular, that is, rows(columns) of arrays in batch should be the same.

One idea is to choose a image of the same size for each stack. Caution.

  1. shuffling: can group indexes/images_id by size and then apply shuffle within the group.
  2. last batch: just do something similar to drop_last if neccessary (see torch dataloader)
  3. maybe there is more work ...
Related