I am training a model to segment 3D images in a slice by slice-fashion. To distribute over multiple GPUs I am using DistributedDataParallel and I use DistributedSampler to split the dataset across the GPUs.
During prediction of new cases, I use a similar Dataset and DataLoader setup and I basically can gather a dictionary like: {'filename': [(slice_no_1, slice_1_pred), (slice_no_2, slice_2_pred)], ...} which I can subsequently sort on the first index to get an output. However, when I use DistributedSampler the slices are distributed along two GPUs, and I therefore end up with two dictionaries which most likely are both incomplete (one containing slices of the other and vice versa).
How do I gather these two dictionaries? As I preferably cast the predictions to a numpy array, it might be most convenient to gather these in the CPU memory.