File size: 2,279 Bytes
032e687 |
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 |
import torch
from torch.utils.data import Sampler, BatchSampler
from typing import Iterator, Optional, Sized
from mmengine.dist import get_dist_info
class MultiDataPseudoSampler(Sampler):
def __init__(self, dataset, seed=None, round_up=True):
self.dataset = dataset
def __iter__(self) -> Iterator[int]:
pass
def __len__(self) -> int:
return self.num_samples
def set_epoch(self, epoch: int) -> None:
self.epoch = epoch
class MultiDataSameBatchSampler(BatchSampler):
def __init__(self, sampler, batch_size: int, drop_last: bool = True):
self.sampler = sampler
self.batch_size = batch_size
self.drop_last = drop_last
rank, world_size = get_dist_info()
self.world_size = world_size
self.rank = rank
self.dataset = sampler.dataset
total_batches = 0
for start, end in zip([0] + self.dataset.cumulative_sizes[:-1], self.dataset.cumulative_sizes):
total_batches_ = (end - start) // self.batch_size
total_batches += total_batches_
self.num_samples = total_batches // self.world_size * self.batch_size
self.total_size = self.num_samples * self.world_size
self.epoch = 0
def __iter__(self):
indices = self._shuffle()
indices = indices[self.rank:self.total_size // self.batch_size:self.world_size]
assert len(indices) * self.batch_size == self.num_samples
return iter(indices)
def _shuffle(self):
g = torch.Generator()
g.manual_seed(42 + self.epoch)
indices = []
for start, end in zip([0] + self.dataset.cumulative_sizes[:-1], self.dataset.cumulative_sizes):
indices_ = torch.randperm(end-start, generator=g) + start
if len(indices_) % self.batch_size:
indices_ = indices_[:-(len(indices_) % self.batch_size)]
indices_ = indices_.view(-1, self.batch_size)
indices += indices_
indices = torch.stack(indices)
indices = indices[torch.randperm(len(indices), generator=g)]
return indices.tolist()
def __len__(self) -> int:
return self.num_samples // self.batch_size
def set_epoch(self, epoch):
self.epoch = epoch
|