File size: 2,886 Bytes
f53b39e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# Random sampling under a constraint
# --------------------------------------------------------
import numpy as np
import torch
class BatchedRandomSampler:
""" 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, world_size=1, rank=0, drop_last=True):
self.batch_size = batch_size
self.pool_size = pool_size
self.len_dataset = N = len(dataset)
self.total_size = round_by(N, batch_size*world_size) if drop_last else N
assert world_size == 1 or drop_last, 'must drop the last batch in distributed mode'
# distributed sampler
self.world_size = world_size
self.rank = rank
self.epoch = None
def __len__(self):
return self.total_size // self.world_size
def set_epoch(self, epoch):
self.epoch = epoch
def __iter__(self):
# prepare RNG
if self.epoch is None:
assert self.world_size == 1 and self.rank == 0, 'use set_epoch() if distributed mode is used'
seed = int(torch.empty((), dtype=torch.int64).random_().item())
else:
seed = self.epoch + 777
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
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]
# put them together
idxs = np.c_[sample_idxs, feat_idxs] # shape = (total_size, 2)
# Distributed sampler: we select a subset of batches
# make sure the slice for each node is aligned with batch_size
size_per_proc = self.batch_size * ((self.total_size + self.world_size *
self.batch_size-1) // (self.world_size * self.batch_size))
idxs = idxs[self.rank*size_per_proc: (self.rank+1)*size_per_proc]
yield from (tuple(idx) for idx in idxs)
def round_by(total, multiple, up=False):
if up:
total = total + multiple-1
return (total//multiple) * multiple
|