Using torchio for DDP

[alabamagan]

I can't seem to find any discussion of using torchio with distributed model. I think the challenge mainly lies at how the dataloader samplers can distribute between processes of different ranks. Does anyone has the idea how this can be achieved with torchio?

[fepegar]

I think @hsyang1222 might be able to help. See also

• Improve support for queue in distributed training #1021

[alabamagan]

Fabulous, I currently just create separate Queue on separate process and you can imagine how messy it would get. I will have a look at the PR.

[hsyang1222]

We recently added feature for support DDP samplers in torchio queues. Please see the queue section of the updated documentation. If you don't understand the documentation, please let me know.

[alabamagan]

Seems like the Queue in each process will loop through the whole data set so that if you are DDP training with a world size N, each epoch will be trained with N x [num_samples]. Is this the default behavior or am I missing something?

[hsyang1222]

As the world size N increases, the number of samples learned in each epoch num_samples remains the same.

First, without using DDP, num_samples would be calculated as follows.

Number of subjects S * number of patch extractions per subject (samples_per_volume) P = total number of samples num_samples

When using a distributed sampler and queue in an environment with a DDP of world size N, the distributed sampler distributes the subjects to be used at each rank. In each rank, num_samples is calculated as follows

num_samples in rank 0 : (S // N) * P
num_samples in rank 1 : (S // N) * P
...
num_samples in rank N-1 : (S // N) * P

The sum of these is (S // N) * P * N, which is the same number as without DDP. Of course, there is some margine when S is not divisible by N, but it is negligible.