How to combine two different samplers? MNIST

username: mgloria
cooked: <p>I am doing distributed training with the mnist dataset. The mnist dataset is only split (by default) between training and testing set. I would like to split the training set in training and validation set.</p>
<p>I could do this as follows:</p>
<pre><code class="lang-auto"># Shuffle the indices
indices = np.arange(60000)
np.random.shuffle(indices)

# Build the train loader
train_loader = torch.utils.data.DataLoader(datasets.MNIST('mnist', download=True, train=True,
                     transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])),
                     batch_size=64, shuffle=False, sampler=torch.utils.data.SubsetRandomSampler(indices[:55000]))

# Build the validation loader
val_loader = torch.utils.data.DataLoader(datasets.MNIST('mnist', download=True, train=True,
                   transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])),
                   batch_size=64, shuffle=False, sampler=torch.utils.data.SubsetRandomSampler(indices[55000:]))
</code></pre>
<p>but how do I incorporate the distributed training? Before (without split) I was doing it as follows:</p>
<pre><code class="lang-auto"># gets data from training-dir (s3 bucket path)
def _get_train_data_loader(batch_size, training_dir, is_distributed, **kwargs):
    logger.info("Get train data loader")
      
    dataset = datasets.MNIST(training_dir, train=True, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ]))
    train_sampler = torch.utils.data.distributed.DistributedSampler(dataset) if is_distributed else None
    return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=train_sampler is None,
                                       sampler=train_sampler, **kwargs)
</code></pre>
<p>The problem is that I do not know how to pass 2 samplers to do both things when creating the dataloader.</p>
<p>I found this <a href="https://github.com/pytorch/pytorch/issues/23430" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/23430</a> which seems to be related but since I am a beginner I am not really able to make sense of it.</p>

username: ptrblck
cooked: <p>The cleanest approach would probably be to create a <code>DistributedSubsetSampler</code>, if you want to combine these both samplers together.<br>
However, I think the easiest workaround would be to use <code>DistributedSampler</code> and wrap your <code>Dataset</code> in <code>Subset</code> using the desired indices.</p>

username: mgloria
cooked: <p>Thanks a lot <a class="mention" href="/u/ptrblck">@ptrblck</a> for your awesome help! I did not know of <code>Subset</code> which totally did the trick in one line. I am adding below the solution for future reference.</p>
<pre><code class="lang-auto">def _get_train_data_loader(batch_size, training_dir, is_distributed, indices, **kwargs):
    logger.info("Get train data loader")
    
    # Train = True, indices[:55000]
    dataset = datasets.MNIST(training_dir, train=True, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ]))
    
    train_ds  = torch.utils.data.Subset(dataset, indices[:55000])
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_ds) if is_distributed else None
    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=batch_size, shuffle=train_sampler is None, sampler=train_sampler, **kwargs)
    return train_loader

def _get_valid_data_loader(valid_batch_size, training_dir, indices, **kwargs):
    logger.info("Get valid data loader")
    
    # train = True, indices[55000:]
    dataset = datasets.MNIST(training_dir, train=True, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ]))
    
    valid_ds  = torch.utils.data.Subset(dataset, indices[55000:])
    valid_loader = torch.utils.data.DataLoader(valid_ds, batch_size=valid_batch_size, shuffle=True, **kwargs)
    return valid_loader

def _get_test_data_loader(test_batch_size, training_dir, **kwargs):
    logger.info("Get test data loader")
    
    # train = False
    dataset = datasets.MNIST(training_dir, train=False, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ]))
    
    test_loader= torch.utils.data.DataLoader(dataset, batch_size=test_batch_size, shuffle=True, **kwargs)
    return test_loader 
</code></pre>
<p>Note that I am calling these function from the train function like this:</p>
<pre><code class="lang-auto">    indices = np.arange(60000)
    np.random.shuffle(indices)    
    train_loader = _get_train_data_loader(args.batch_size, args.data_dir, is_distributed, indices, **kwargs)
    valid_loader = _get_valid_data_loader(args.valid_batch_size, args.data_dir, indices, **kwargs)
    test_loader = _get_test_data_loader(args.test_batch_size, args.data_dir, **kwargs)```</code></pre>

username: Ammara
cooked: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a><br>
How will I create the <code>DistributedSubsetSampler</code>?<br>
I have to combine <code>WeightedRandomSampler</code> and <code>DistributedSampler</code> for batch balancing with ddp training.</p>

username: ptrblck
cooked: <p><a href="https://discuss.pytorch.org/t/weightedrandomsampler-distributedsampler/52817/5">This post</a> and thread discusses the use case of combining both samplers so you could take a look at some of the posted implementations.</p>