Spaces:
Running
Running
File size: 4,585 Bytes
64bf706 |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
import numpy as np
import torch
from torch.utils.data.sampler import Sampler
class EvalDistributedSampler(Sampler):
def __init__(self, dataset, num_replicas, rank):
seps = np.linspace(0, len(dataset), num_replicas+1, dtype=int)
beg, end = seps[:-1], seps[1:]
beg, end = beg[rank], end[rank]
self.indices = tuple(range(beg, end))
def __iter__(self):
return iter(self.indices)
def __len__(self) -> int:
return len(self.indices)
class InfiniteBatchSampler(Sampler):
def __init__(self, dataset_len, batch_size, seed_for_all_rank=0, fill_last=False, shuffle=True, drop_last=False, start_ep=0, start_it=0):
self.dataset_len = dataset_len
self.batch_size = batch_size
self.iters_per_ep = dataset_len // batch_size if drop_last else (dataset_len + batch_size - 1) // batch_size
self.max_p = self.iters_per_ep * batch_size
self.fill_last = fill_last
self.shuffle = shuffle
self.epoch = start_ep
self.same_seed_for_all_ranks = seed_for_all_rank
self.indices = self.gener_indices()
self.start_ep, self.start_it = start_ep, start_it
def gener_indices(self):
if self.shuffle:
g = torch.Generator()
g.manual_seed(self.epoch + self.same_seed_for_all_ranks)
indices = torch.randperm(self.dataset_len, generator=g).numpy()
else:
indices = torch.arange(self.dataset_len).numpy()
tails = self.batch_size - (self.dataset_len % self.batch_size)
if tails != self.batch_size and self.fill_last:
tails = indices[:tails]
np.random.shuffle(indices)
indices = np.concatenate((indices, tails))
# built-in list/tuple is faster than np.ndarray (when collating the data via a for-loop)
# noinspection PyTypeChecker
return tuple(indices.tolist())
def __iter__(self):
self.epoch = self.start_ep
while True:
self.epoch += 1
p = (self.start_it * self.batch_size) if self.epoch == self.start_ep else 0
while p < self.max_p:
q = p + self.batch_size
yield self.indices[p:q]
p = q
if self.shuffle:
self.indices = self.gener_indices()
def __len__(self):
return self.iters_per_ep
class DistInfiniteBatchSampler(InfiniteBatchSampler):
def __init__(self, world_size, rank, dataset_len, glb_batch_size, same_seed_for_all_ranks=0, repeated_aug=0, fill_last=False, shuffle=True, start_ep=0, start_it=0):
assert glb_batch_size % world_size == 0
self.world_size, self.rank = world_size, rank
self.dataset_len = dataset_len
self.glb_batch_size = glb_batch_size
self.batch_size = glb_batch_size // world_size
self.iters_per_ep = (dataset_len + glb_batch_size - 1) // glb_batch_size
self.fill_last = fill_last
self.shuffle = shuffle
self.repeated_aug = repeated_aug
self.epoch = start_ep
self.same_seed_for_all_ranks = same_seed_for_all_ranks
self.indices = self.gener_indices()
self.start_ep, self.start_it = start_ep, start_it
def gener_indices(self):
global_max_p = self.iters_per_ep * self.glb_batch_size # global_max_p % world_size must be 0 cuz glb_batch_size % world_size == 0
# print(f'global_max_p = iters_per_ep({self.iters_per_ep}) * glb_batch_size({self.glb_batch_size}) = {global_max_p}')
if self.shuffle:
g = torch.Generator()
g.manual_seed(self.epoch + self.same_seed_for_all_ranks)
global_indices = torch.randperm(self.dataset_len, generator=g)
if self.repeated_aug > 1:
global_indices = global_indices[:(self.dataset_len + self.repeated_aug - 1) // self.repeated_aug].repeat_interleave(self.repeated_aug, dim=0)[:global_max_p]
else:
global_indices = torch.arange(self.dataset_len)
filling = global_max_p - global_indices.shape[0]
if filling > 0 and self.fill_last:
global_indices = torch.cat((global_indices, global_indices[:filling]))
# global_indices = tuple(global_indices.numpy().tolist())
seps = torch.linspace(0, global_indices.shape[0], self.world_size + 1, dtype=torch.int)
local_indices = global_indices[seps[self.rank].item():seps[self.rank + 1].item()].tolist()
self.max_p = len(local_indices)
return local_indices
|