DynamicBatchSampler

[wangleiofficial]

I adopt my own defined implementation of DistributedSampler and BatchSampler.

from torch.utils.data import Dataset,  DistributedSampler，BatchSampler
class DynamicBatchSampler(BatchSampler):
    """
	Parameters:
	-----------
	sampler : Pytorch Sampler
		Choose base sampler class to use for bucketing
	max_tokens : int
		Maximum number of tokens per batch
	max_batch: int
		Maximum batch size
	sample_lengths : array-like
		List of lengths of sequences in the order of the dataset
	"""

    def __init__(self, sampler, drop_last, max_tokens, max_batch, sample_lengths, max_square_tokens=np.inf):
        self.longest_token = 0
        self.max_tokens = max_tokens
        self.max_batch = max_batch
        self.sampler = sampler
        self.sample_lengths = sample_lengths
        self.max_square_tokens = max_square_tokens
        self.drop_last = drop_last

    def __iter__(self):
        batch = []
        length = 0
        ell_sq = 0
        for idx in self.sampler:
            this_length = self.sample_lengths[idx]
            linear = (len(batch) + 1) * max(length, this_length)
            quadratic = (len(batch) + 1) * max(ell_sq, this_length ** 2)
            if linear <= self.max_tokens and quadratic < self.max_square_tokens:
                batch.append(idx)
                length = max(length, this_length)
                ell_sq = max(ell_sq, this_length ** 2)
                if len(batch) == self.max_batch:
                    yield batch
                    batch = []
                    length = 0
            else:
                yield batch
                batch = [idx]
                length = this_length
                ell_sq = this_length ** 2
        if len(batch) > 0:
            yield batch

pytorch lighting needs to implement len function, the dynamic batch size is difficult to obtain, how to implement？