def setup(self, stage=None):
mb_train = MB_nohybrids(self.path,
train=True,
transform=transforms["mb"])
mb_test = MB_nohybrids(self.path,
train=False,
transform=test_transforms["mb"])
self.test_dataset = mb_test
if self.fraction != 1:
mb_train, _ = d_split(mb_train, self.fraction)
self.train_dataset = mb_train
mb_l, mb_u = d_split(mb_train, self.split)
# Flip a number of targets randomly
if config["train"]["flip"]:
mb_u = flip_targets(mb_u, config["train"]["flip"])
self.train_dataset_u = mb_u
self.train_dataset_l = mb_l
self.f_u = label_fraction(mb_u, 0)
self.f_l = label_fraction(mb_l, 0)
def save_hparams(self):
self.hparams.update({
"n_labelled": len(self.train_l),
"n_unlabelled": len(self.train_u),
"n_test": len(self.test),
"f_u": label_fraction(self.train_u, 0),
"f_l": label_fraction(self.train_l, 0),
})