Repository: pytorch
Issue: BucketSampler for easy variable-length input batching
Labels: feature,module: dataloader,triaged

## 🚀 Feature
An additional [Sampler](https://pytorch.org/docs/stable/data.html#torch.utils.data.Sampler) that helps dealing with generating batches of the same (or close) length in order to avoid padding or packing.

## Motivation

In many different scenarios involving input (or output) that contains sequences, there's the need to deal with the fact that not all sequences in the dataset will be the exact same length and therefore one needs to pad or pack the sequence.
Padding is generally not efficient and packing may be hard to implement in more complex model architectures.
Bucketing samples that are close in length into the same "length bucket" is both efficient and the architecture is agnostic to this since all the batch is of the same length.

## Pitch

Using the `batch_sampler` argument of `DataLoader` it's very easy to generate batches from the same length bucket.
I suggest adding another `Sampler` to the existing ones (`Random`, `Sequential`, etc.) called `BucketSampler` using the following scheme (that code works functionally but obviously is not ideal):

```
import random
from collections import defaultdict

class BucketSampler(torch.utils.data.Sampler):
    
    def __init__(self, lengths, buckets=(50,500,50), shuffle=True, batch_size=32, drop_last=False):
        
        super().__init__(lengths)
        
        self.shuffle = shuffle
        self.batch_size = batch_size
        self.drop_last = drop_last
        
        assert isinstance(buckets, tuple)
        bmin, bmax, bstep = buckets
        assert (bmax - bmin) % bstep == 0
        
        buckets = defaultdict(list)
        for i, length in enumerate(lengths):
            if length > bmin:
                bucket_size = min((length // bstep) * bstep, bmax)
                buckets[bucket_size].append(i)
                
        self.buckets = dict()
        for bucket_size, bucket in buckets.items():
            if len(bucket) > 0:
                self.buckets[bucket_size] = torch.tensor(bucket, dtype=torch.int, device='cpu')
        
        # call __iter__() to store self.length
        self.__iter__()
            
    def __iter__(self):
        
        if self.shuffle == True:
            for bucket_size in self.buckets.keys():
                self.buckets[bucket_size] = self.buckets[bucket_size][torch.randperm(self.buckets[bucket_size].nelement())]
                
        batches = []
        for bucket in self.buckets.values():
            curr_bucket = torch.split(bucket, self.batch_size)
            if len(curr_bucket) > 1 and self.drop_last == True:
                if len(curr_bucket[-1]) < len(curr_bucket[-2]):
                    curr_bucket = curr_bucket[:-1]
            batches += curr_bucket
            
        self.length = len(batches)
        
        if self.shuffle == True:
            random.shuffle(batches)
            
        return iter(batches)
    
    def __len__(self):
        return self.length
```

If this `BucketSampler` exists in `torch.utils.data.BucketSampler`, then all the user needs to do in order to have the data loader generate appropriate batches is:

```
dataset = torch.utils.data.Dataset(...)
bucket_sampler = torch.utils.data.BucketSampler(lengths, batch_size=32, shuffle=True, drop_last=False)
loader = torch.utils.data.DataLoader(dataset, batch_sampler=bucket_sampler)
```

Calculating the lengths of all the sequences in the dataset is pretty straight forward and can be done once and stored.

## Additional context

My suggestion draws from several other methods of handling this use case such as: https://github.com/pytorch/pytorch/issues/25743, https://github.com/pytorch/vision/blob/d5379656a098f0d69ad4dbe3dd94f2701415824c/references/detection/group_by_aspect_ratio.py#L23, and comments in https://github.com/pytorch/pytorch/issues/41292.

Since `torchtext` and `torchvision` already have similar classes, it'll be a good and very practical addition to the `DataLoader`'s capabilities.

cc @SsnL @VitalyFedyunin @ejguan