# Using DistributedSampler in combination with batch\_sampler to make sure batches have sentences of similar length

**URL:** <https://discuss.pytorch.org/t/using-distributedsampler-in-combination-with-batch-sampler-to-make-sure-batches-have-sentences-of-similar-length/119824>\
**Category:** nlp\
**Created:** [April 29, 2021, 1:59pm UTC](https://discuss.pytorch.org/t/using-distributedsampler-in-combination-with-batch-sampler-to-make-sure-batches-have-sentences-of-similar-length/119824 "2021-04-29T13:59:03Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![nbortych](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/nbortych/32/33816_2.png) [@nbortych](https://discuss.pytorch.org/u/nbortych)\
**Post date:** [April 29, 2021, 1:59pm UTC](https://discuss.pytorch.org/t/using-distributedsampler-in-combination-with-batch-sampler-to-make-sure-batches-have-sentences-of-similar-length/119824/1 "2021-04-29T13:59:03Z")

</div>

With torchtext 0.9.0, BucketIterator was depreciated and DataLoader is encouraged to be used instead, which is great since DataLoader is compatible with [DistributedSampler](https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler) and hence DDP. However, it has a downside of not having the out-of-the-box implementation of having batches of similar length. The [migration tutorial](https://colab.research.google.com/github/zhangguanheng66/text/blob/migration_tutorial/examples/legacy_tutorial/migration_tutorial.ipynb#scrollTo=zCvxeLbYW3I_) recommends using batch\_sampler argument of DataLoader to pool together batches of similar length. Unfortunately, the batch\_sampler is not compatible with the sampler.

Does anyone have any suggestions/ideas on how to make sure that when using DDP, the batches are of as similar length as possible?

---

<div class="post-metadata">

**Author:** ![nbortych](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/nbortych/32/33816_2.png) [@nbortych](https://discuss.pytorch.org/u/nbortych)\
**Post date:** [May 27, 2021, 8:10am UTC](https://discuss.pytorch.org/t/using-distributedsampler-in-combination-with-batch-sampler-to-make-sure-batches-have-sentences-of-similar-length/119824/3 "2021-05-27T08:10:24Z")

</div>

I’ve solved the problem by hacking together the [Distributed sampler](https://pytorch.org/docs/stable/_modules/torch/utils/data/distributed.html#DistributedSampler) and the batch\_sampler function from the [migration tutorial](https://colab.research.google.com/github/zhangguanheng66/text/blob/migration_tutorial/examples/legacy_tutorial/migration_tutorial.ipynb#scrollTo=Jky4X-iFU4HK).  
Note that this implementation is for a seq2seq problem and the dataset is a tuple with 4 values: source seq, length of source seq, target seq and length of the target seq. It’s not important at all after the first lines.  
I’ve abstracted the length sorting into a BatchSamplerSimilarLength Sampler class which can be used without any DDP.

```auto
class BatchSamplerSimilarLength(Sampler):
    def __init__ (self, dataset, batch_size, indices=None, shuffle=True):
        self.batch_size = batch_size
        self.shuffle = shuffle
        # get the indicies and length
        self.indices = [(i, src_len) for i, (src, src_len, trg, trg_len) in enumerate(dataset)]
        # if indices are passed, then use only the ones passed (for ddp)
        if indices is not None:
            self.indices = torch.tensor(self.indices)[indices].tolist()

    def __iter__ (self):
        if self.shuffle:
            random.shuffle(self.indices)

        pooled_indices = []
        # create pool of indices with similar lengths
        for i in range(0, len(self.indices), self.batch_size * 100):
            pooled_indices.extend(sorted(self.indices[i:i + self.batch_size * 100], key=lambda x: x[1]))
        self.pooled_indices = [x[0] for x in pooled_indices]

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

        if self.shuffle:
            random.shuffle(batches)
        for batch in batches:
            yield batch

    def __len__ (self):
        return len(self.pooled_indices) // self.batch_size

```

Then I’ve created a child of the Distributed Sampler that overrides the **iter** method by passing the indices computed by the Distributed Sampler to the BatchSampler, so the Batch Sampler has access only to the indices at the current rank and then can proceed as normal (by batching by length).

```auto
class DistributedBatchSamplerSimilarLength(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, batch_size = 10) -> None:
        super(). __init__ (dataset=dataset, num_replicas=num_replicas, rank=rank, shuffle=shuffle, seed=seed,
                         drop_last=drop_last)
        self.batch_size = batch_size

    def __iter__ (self):
        indices = list(super(). __iter__ ())
        batch_sampler = BatchSamplerSimilarLength(self.dataset, batch_size=self.batch_size, indices=indices)
        return iter(batch_sampler)

    def __len__ (self) -> int:
        return self.num_samples//batch_size

```

The DistributedBatchSamplerSimilarLength has to be passed as a `batch_sampler` argument to the DataLoader.

---

<div class="post-metadata">

**Author:** ![andrei.radulescubanu](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/andrei.radulescubanu/32/55605_2.png) [@andrei.radulescubanu](https://discuss.pytorch.org/u/andrei.radulescubanu)\
**Post date:** [December 10, 2022, 11:43pm UTC](https://discuss.pytorch.org/t/using-distributedsampler-in-combination-with-batch-sampler-to-make-sure-batches-have-sentences-of-similar-length/119824/4 "2022-12-10T23:43:48Z")

</div>

Terrific. I spent a day with this before running into your solution. It works!
