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}
