How to implement a custom distributed sampler

mk6: <p>Hi,</p>
<p>I’m working on sequence data and would like to group sequences of similar lengths into batches. Here is an example implementation (<a href="https://colab.research.google.com/github/pytorch/text/blob/master/examples/legacy_tutorial/migration_tutorial.ipynb#scrollTo=zCvxeLbYW3I_&amp;line=24&amp;uniqifier=1" rel="noopener nofollow ugc">source</a>)</p>
<pre><code class="lang-auto">"""
To group the texts with similar length together, like introduced in the legacy BucketIterator class,
first of all, we randomly create multiple "pools", and each of them has a size of batch_size * 100.
Then, we sort the samples within the individual pool by length. 
This idea can be implemented succintly through batch_sampler argument of PyTorch Dataloader. 
batch_sampler accepts 'Sampler' or Iterable object that yields indices of next batch. 
In the code below, we implemented a generator that yields batch of indices for which the corresponding batch of data is of similar length.
"""
import random

train_iter = IMDB(split='train')
train_list = list(train_iter)
batch_size = 8  # A batch size of 8

def batch_sampler():
    indices = [( i, len(tokenizer(s[1])) ) for i, s in enumerate(train_list)]
    random.shuffle(indices)
    pooled_indices = []
    # create pool of indices with similar lengths 
    for i in range(0, len(indices), batch_size * 100):
        pooled_indices.extend(sorted(indices[i:i + batch_size * 100], key=lambda x: x[1]))

    pooled_indices = [x[0] for x in pooled_indices]

    # yield indices for current batch
    for i in range(0, len(pooled_indices), batch_size):
        yield pooled_indices[i:i + batch_size]

bucket_dataloader = DataLoader(train_list, batch_sampler=batch_sampler(),
                               collate_fn=collate_batch)

</code></pre>
<p>I’m wondering what should I do to adapt it to a distributed one.</p>
<p>Thanks</p>

Matias_Vasquez: <p>What you can do is create a class that inherits from <code>DistributedSampler</code> and modify it so that it behaves to what you need.</p>
<pre><code class="lang-python">from torch.utils.data.distributed import DistributedSampler

class MyDistributedBatchSampler(DistributedSampler):
    def __init__(self, dataset: Dataset, num_replicas: Optional[int] = None,
                 rank: Optional[int] = None, shuffle: bool = True,
                 seed: int = 0, drop_last: bool = False) -&gt; None:
        super().__init__(self, dataset, num_replicas, rank, shuffle,
                         seed, drop_last)
        
        # Initialize your stuff

    def __iter__(self) -&gt; Iterator[List[int]]:
        # Here you need to look how DistributedSampler implements __iter__
        # https://pytorch.org/docs/stable/_modules/torch/utils/data/distributed.html#DistributedSampler
        # Then you can modify it to serve your purposes

    # If len stays the same you can leave it out, else you can also modify it
    #def __len__(self):
        #return self.num_samples
</code></pre>

mk6: <p>Hi Matias,</p>
<p>Thanks for help.<br>
I just found I shouldn’t implement my own <code>Sampler</code> since I’m dealing with <code>Iterable</code> datasets, as suggested in the doc <a href="https://pytorch.org/docs/stable/data.html#data-loading-order-and-sampler" rel="noopener nofollow ugc">here</a>.</p>

Ryan_Wang: <p>just adding on top of that. I recently realized that we have to derive from DistributedSampler here, rather than normal sampler, or data sharding will happen 2 times in the mutl-gpu training.</p>