Sample one batch from a selected dataset with given probability. Compatible with datasets at different resolution
| 29 | |
| 30 | |
| 31 | class MixedBatchSampler(BatchSampler): |
| 32 | """Sample one batch from a selected dataset with given probability. |
| 33 | Compatible with datasets at different resolution |
| 34 | """ |
| 35 | |
| 36 | def __init__( |
| 37 | self, src_dataset_ls, batch_size, drop_last, shuffle, prob=None, generator=None |
| 38 | ): |
| 39 | self.base_sampler = None |
| 40 | self.batch_size = batch_size |
| 41 | self.shuffle = shuffle |
| 42 | self.drop_last = drop_last |
| 43 | self.generator = generator |
| 44 | |
| 45 | self.src_dataset_ls = src_dataset_ls |
| 46 | self.n_dataset = len(self.src_dataset_ls) |
| 47 | |
| 48 | # Dataset length |
| 49 | self.dataset_length = [len(ds) for ds in self.src_dataset_ls] |
| 50 | self.cum_dataset_length = [ |
| 51 | sum(self.dataset_length[:i]) for i in range(self.n_dataset) |
| 52 | ] # cumulative dataset length |
| 53 | |
| 54 | # BatchSamplers for each source dataset |
| 55 | if self.shuffle: |
| 56 | self.src_batch_samplers = [ |
| 57 | BatchSampler( |
| 58 | sampler=RandomSampler( |
| 59 | ds, replacement=False, generator=self.generator |
| 60 | ), |
| 61 | batch_size=self.batch_size, |
| 62 | drop_last=self.drop_last, |
| 63 | ) |
| 64 | for ds in self.src_dataset_ls |
| 65 | ] |
| 66 | else: |
| 67 | self.src_batch_samplers = [ |
| 68 | BatchSampler( |
| 69 | sampler=SequentialSampler(ds), |
| 70 | batch_size=self.batch_size, |
| 71 | drop_last=self.drop_last, |
| 72 | ) |
| 73 | for ds in self.src_dataset_ls |
| 74 | ] |
| 75 | self.raw_batches = [ |
| 76 | list(bs) for bs in self.src_batch_samplers |
| 77 | ] # index in original dataset |
| 78 | self.n_batches = [len(b) for b in self.raw_batches] |
| 79 | self.n_total_batch = sum(self.n_batches) |
| 80 | # sampling probability |
| 81 | if prob is None: |
| 82 | # if not given, decide by dataset length |
| 83 | self.prob = torch.tensor(self.n_batches) / self.n_total_batch |
| 84 | else: |
| 85 | self.prob = torch.as_tensor(prob) |
| 86 | |
| 87 | def __iter__(self): |
| 88 | """_summary_ |