ptrblck:

Given that each batch has the shape [batch_size, subbatch_size, channels, x_resolution, y_resolution], are the sizes in each dimension equal for each batch or would e.g. subbatch_size be different?


batch shapes will be the same across all bachtes durint training.



 ptrblck:

In the latter case, you could either flatten the subbatch patches into the batch dimension and use the default collate_fn or, if that’s not desired, you could write a custom collate_fn.
However, this also depends on the expected model input. I.e. would it work, if you put all patches into the batch dimension or do you want to process these patches somehow differently?


Yea, you made me realize that i actually dont need to get the data in the shape [batch_size, subbatch_size, channels, x_resolution, y_resolution]. [batch_size*subbatch_size, channels, x_resolution, y_resolution] would be fine, since my model will consume the batches like that anyway (it acts only on the tiles). Just my loss and maybe some other metrics depend on the subbatched predictions. So flattening and using the default collate_fn would be a good approach as well.
But:
After some testing, i made it work by making the __get_item__ method of my dataset accept the subbatched indices of my Sampler and return subbatches. I guess i might sacrifice a tiny bit of performance this way when batch_size % num_workers != 0, as now each worker will collect subbatches instead of individual samples? But i assume it will be no real issue, since the dataloader does prefetching anyway…
This way i don’t need an additional custom BatchSampler and also can inferre the subbatch_size from the data input to my model directly and dont need to pass it as an argument. Since my model needs to be able to accept any (or no) subbatch_size after training, that’s my preffered way.
So now all i need to do is distributify my sampler and im good to go!