Multiple dataloaders to be able to combine tiles of large images

[jonasteuwen]

I'm working with very large images, and these are typically processed in tiles. The final prediction is an aggregation of tiles. For instance, if I have an image of 25,000 x 25,000 pixels, I might split it up in 100 images of size 2500x2500. Memory limitations aside, I would like to do the evaluation on the full image and not the tiles. So I need to aggregate the results per image.

I am wondering what would be the most elegant way to do this in pytorch lightning. What I currently have is that each image itself is a dataset of tiles, and I use ConcatDataset to combine all of these during training. However, I saw the possibility to use multiple dataloaders during validation by returning a list of those. If that would work, I suppose I could build a custom EvaluationLoop that combines the results at the end of an dataloader.

However, that's not what happens, and in that case lightning returns a list of samples. Ideally, I would also want to have a generator returning a new Dataloader each time, so like

def _shared_dataloader(self):
    for dataset in self.datasets:
      yield DataLoader(
          self.datasets[prefix],
          batch_size=self.batch_size,
          pin_memory=True,
          collate_fn=custom_collate_fn,
          shuffle=True,
          num_workers=self.num_workers,
      )

as I can have thousands of images. What would be the way to approach this in Lightning?

[NicoMandel]

Have you found a solution to this, I am currently looking at ways to solve this elegantly