def setup_train_dataset(self): """ Each self.batch_size of examples follows the same distribution """ bd = BlockDesigner(self.train_examples) if self.sample_class: samp = Sampler(bd.remainder(), seed=self.random_seed) images, labels = samp.custom_distribution(self.sample_class, self.batch_size, self.custom_distribution) return {"X": images, "y": labels} else: blocks = bd.break_off_multiple_blocks(self.n_train_batches, self.batch_size) images = [] labels = [] for block in blocks: for y, ids in block.items(): for id in ids: images.append(id) labels.append(y) return {"X": images, "y": labels}