Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import torch | |
| from accelerate import Accelerator | |
| import torch.utils | |
| from torch.utils.data import BatchSampler, Sampler | |
| import torch.utils.data | |
| class CustomRandomSampler(Sampler): | |
| """Random sampling under a constraint: each sample in the batch has the same feature, | |
| which is chosen randomly from a known pool of 'features' for each batch. | |
| For instance, the 'feature' could be the image aspect-ratio. | |
| The index returned is a tuple (sample_idx, feat_idx). | |
| This sampler ensures that each series of `batch_size` indices has the same `feat_idx`. | |
| """ | |
| def __init__( | |
| self, | |
| dataset, | |
| batch_size, | |
| pool_size, | |
| min_view_size, | |
| max_view_size, | |
| world_size, | |
| warmup=1, | |
| drop_last=True, | |
| ): | |
| self.batch_size = batch_size | |
| self.pool_size = pool_size | |
| self.min_view_size = min_view_size | |
| self.max_view_size = max_view_size | |
| self.drop_last = drop_last | |
| self.len_dataset = N = len(dataset) | |
| self.total_size = N | |
| self.epoch = None | |
| self.epochf = 0.0 | |
| def __len__(self): | |
| return self.total_size | |
| def set_epoch(self, epoch): | |
| self.epoch = epoch | |
| def __iter__(self): | |
| if self.epoch is None: | |
| raise ValueError( | |
| "Epoch number not set. Please call 'set_epoch(epoch)' before iterating." | |
| ) | |
| seed = self.epoch + 788 | |
| rng = np.random.default_rng(seed=seed) | |
| # random indices (will restart from 0 if not drop_last) | |
| sample_idxs = np.arange(self.total_size) | |
| rng.shuffle(sample_idxs) | |
| # random feat_idxs (same across each batch) | |
| n_batches = (self.total_size + self.batch_size - 1) // self.batch_size | |
| if self.pool_size > 1: | |
| p = np.ones(self.pool_size) | |
| p[: self.pool_size // 2] *= 2 | |
| p = p / p.sum() | |
| _feat_idxs = rng.choice(self.pool_size, size=n_batches, p=p) | |
| else: | |
| _feat_idxs = rng.integers(self.pool_size, size=n_batches) | |
| _feat_idxs = np.broadcast_to(_feat_idxs[:, None], (n_batches, self.batch_size)) | |
| _feat_idxs = _feat_idxs.ravel()[: self.total_size] | |
| _view_idxs = rng.integers( | |
| self.min_view_size, self.max_view_size + 1, size=n_batches | |
| ) | |
| _view_idxs = np.broadcast_to(_view_idxs[:, None], (n_batches, self.batch_size)) | |
| _view_idxs = _view_idxs.ravel()[: self.total_size] | |
| idxs = np.c_[sample_idxs, _feat_idxs, _view_idxs] | |
| yield from (tuple(idx) for idx in idxs) | |
| class BatchedRandomSampler(BatchSampler): | |
| """Batch sampler that groups indices from RandomSampler into batches.""" | |
| def __init__(self, sampler: CustomRandomSampler, batch_size, drop_last=True): | |
| self.sampler = sampler # An instance of RandomSampler | |
| self.batch_size = batch_size | |
| self.drop_last = drop_last | |
| def set_epoch(self, epoch): | |
| self.sampler.set_epoch(epoch) | |
| def round_by(total, multiple, up=False): | |
| if up: | |
| total = total + multiple - 1 | |
| return (total // multiple) * multiple | |