steffenN:

What else do i need to do to make this work? Does this just work out of the box when i pass the sampler to the DataLoader? I couldn’t try yet, but i guess not? Do i have to build a custom BatchSampler as well, where i flatten the batch and subbatch dimension into one larger batch dim and reshape my batches back into my wanted shape further down the line?
Or is that whats the collate_fn can be used for, which is passed to the DataLoader?


I think it depends on the actual shape of each batch.
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?
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?
The DistributedSampler is used to split the dataset into chunks, so that each process would only load its subset. Once your custom sampler works fine, you could adapt the split logic from the DistributedSampler.