codingutf8 2019present the huggingface inc team and facebook inc licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license adapted from pytorch vision https github compytorchvisionblobmasterreferencesdetectiongroupbyaspectratio py count number of elements per group wraps another sampler to yield a minibatch of indices it enforces that the batch only contain elements from the same group it also tries to provide minibatches which follows an ordering which is as close as possible to the ordering from the original sampler arguments sampler sampler base sampler groupids listint if the sampler produces indices in range 0 n groupids must be a list of n ints which contains the group id of each sample the group ids must be a continuous set of integers starting from 0 i e they must be in the range 0 numgroups batchsize int size of minibatch now we have run out of elements that satisfy the group criteria let s return the remaining elements so that the size of the sampler is deterministic for the remaining batches group the batches by similar lengths return the number of minibatches rather than the number of samples coding utf 8 2019 present the huggingface inc team and facebook inc licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license adapted from pytorch vision https github com pytorch vision blob master references detection group_by_aspect_ratio py count number of elements per group wraps another sampler to yield a mini batch of indices it enforces that the batch only contain elements from the same group it also tries to provide mini batches which follows an ordering which is as close as possible to the ordering from the original sampler arguments sampler sampler base sampler group_ids list int if the sampler produces indices in range 0 n group_ids must be a list of n ints which contains the group id of each sample the group ids must be a continuous set of integers starting from 0 i e they must be in the range 0 num_groups batch_size int size of mini batch todo now we have run out of elements that satisfy the group criteria let s return the remaining elements so that the size of the sampler is deterministic for the remaining batches group the batches by similar lengths return the number of mini batches rather than the number of samples
import bisect
import copy
from collections import defaultdict

import numpy as np
from torch.utils.data import BatchSampler, Sampler

from utils import logger


def _quantize(x, bins):
    bins = copy.deepcopy(bins)
    bins = sorted(bins)
    quantized = [bisect.bisect_right(bins, y) for y in x]
    return quantized


def create_lengths_groups(lengths, k=0):
    bins = np.arange(start=3, stop=k, step=4).tolist() if k > 0 else [10]
    groups = _quantize(lengths, bins)
    
    counts = np.unique(groups, return_counts=True)[1]
    fbins = [0] + bins + [np.inf]
    logger.info("Using {} as bins for aspect lengths quantization".format(fbins))
    logger.info("Count of instances per bin: {}".format(counts))
    return groups


class GroupedBatchSampler(BatchSampler):
    

    def __init__(self, sampler, group_ids, batch_size):
        if not isinstance(sampler, Sampler):
            raise ValueError(
                "sampler should be an instance of torch.utils.data.Sampler, but got sampler={}".format(sampler)
            )
        self.sampler = sampler
        self.group_ids = group_ids
        self.batch_size = batch_size

    def __iter__(self):
        buffer_per_group = defaultdict(list)
        samples_per_group = defaultdict(list)

        num_batches = 0
        for idx in self.sampler:
            group_id = self.group_ids[idx]
            buffer_per_group[group_id].append(idx)
            samples_per_group[group_id].append(idx)
            if len(buffer_per_group[group_id]) == self.batch_size:
                yield buffer_per_group[group_id]  
                num_batches += 1
                del buffer_per_group[group_id]
            assert len(buffer_per_group[group_id]) < self.batch_size

        
        
        
        
        expected_num_batches = len(self)
        num_remaining = expected_num_batches - num_batches
        if num_remaining > 0:
            
            batch_idx = []
            for group_id, idxs in sorted(buffer_per_group.items(), key=lambda x: x[0]):
                batch_idx.extend(idxs)
                if len(batch_idx) >= self.batch_size:
                    yield batch_idx[: self.batch_size]
                    batch_idx = batch_idx[self.batch_size :]
                    num_remaining -= 1
            if len(batch_idx) > 0:
                yield batch_idx
                num_remaining -= 1
        assert num_remaining == 0

    def __len__(self):
        
        return (len(self.sampler) + self.batch_size - 1) // self.batch_size