How to use pre-batched data with DistributedDataParallel

vindy: <p>Hi all,</p>
<p><strong>Use case</strong>:<br>
I am training a model with image data for a semantic segmentation task. I have a requirement to batch data using ratios of combinations of classes. For example: my batch should have 75% of patches containing classes 0 1 and 2, and 25% of classes containing 0 and 1. Hence, I pre-batched my data.</p>
<p>I want to use this pre-batched data with <code>DistributedDataParallel</code></p>
<p><strong>My question</strong>:<br>
Let’s say <code>n_gpus=4</code>, <code>batch_size=256</code>, and <code>n_batches=10</code>. I understand that if I use the following code, and give <code>num_batches</code> in <code>Dataloader</code> as <code>256</code> It will feed model in each GPU with <code>256/n_gpus=64</code> batch of data. Since I want to ensure I have pre-batched data fed into each model, let’s say I give <code>num_batches</code> in the <code>Dataloader</code> as <code>256*n_gpus=1024</code>, it will do what I want. However, as my n_batches is 10, it won’t equally divide into my GPUs. In the first go, it will take 4 batches and in the second one, it will take another 4. But in the last one, there will be 2 batches left.  My batch size is also dynamic.</p>
<p>How should I handle such a case?</p>
<ul>
<li>
<p>Option 1: Drop remaining batches. Due to my prerequisite in pre-batching, I already drop lots of data. So I do not want to do this.</p>
</li>
<li>
<p>Option 2: Make sure my batch size will always result in <code>n_batches%n_gpu = 0</code>. This is hard (and not optimal) because I have to handle the initial prerequisite of adhering to a certain ratio during the pre-batching.</p>
</li>
<li>
<p>Option 3: Train on 4 GPUs in the first two rounds and use only 2 GPUs in the final round. (Is this possible?)</p>
</li>
</ul>
<p>or else, does Pytorch automatically handle it? (I’m new to Pytorch)</p>
<p><strong>The code</strong>:</p>
<pre><code class="lang-auto">def prepare_dataloader(dataset, current_gpu_index, num_gpus, batch_size):
    sampler = torch.utils.data.distributed.DistributedSampler(
        dataset,
        num_replicas=num_gpus,
        rank=current_gpu_index,
        shuffle=False,
    )
    dataloader = torch.utils.data.DataLoader(
        dataset,
        sampler=sampler,
        batch_size=batch_size,
        shuffle=False,
    )
    return dataloader

def per_device_launch_fn(current_gpu_index, num_gpu):
    # Setup the process groups
    setup_device(current_gpu_index, num_gpu)

    dataset = get_dataset()
    model = get_model()

    # prepare the dataloader
    dataloader = prepare_dataloader(dataset, current_gpu_index, num_gpu, batch_size)

    # Instantiate the torch optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    # Instantiate the torch loss function
    loss_fn = torch.nn.CrossEntropyLoss()

    # Put model on device
    model = model.to(current_gpu_index)
    ddp_model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[current_gpu_index], output_device=current_gpu_index
    )

    train_model(ddp_model, dataloader, num_epochs, optimizer, loss_fn)

    cleanup()

torch.multiprocessing.start_processes(
        per_device_launch_fn,
        args=(num_gpu,),
        nprocs=num_gpu,
        join=True,
        start_method="fork",
    )
</code></pre>

ptrblck: <p>The <code>DistributedSampler</code> commonly used in <code>DDP</code> will duplicate samples to make them evenly divisible as seen <a href="https://github.com/pytorch/pytorch/blob/63c24f73ef53cc7994b7e25b4cdf6b761e8f04bf/torch/utils/data/distributed.py#L106-L111">here</a>. I don’t know if this would fit your requirement, but might be a valid approach instead of dropping samples.</p>

vindy: <p>Yes. It is better than dropping the samples. Thank you for letting me know.</p>