def setup(self, stage: Optional[str] = None):
if stage == "fit" or stage is None:
mnist_full = TrialMNIST(root=self.data_dir,
train=True,
num_samples=64,
download=True)
self.mnist_train, self.mnist_val = random_split(
mnist_full, [128, 64])
self.dims = self.mnist_train[0][0].shape
if stage == "test" or stage is None:
self.mnist_test = TrialMNIST(root=self.data_dir,
train=False,
num_samples=64,
download=True)
self.dims = getattr(self, "dims", self.mnist_test[0][0].shape)
self.non_picklable = lambda x: x**2
def dataloader(self, train: bool, num_samples: int = 100):
dataset = TrialMNIST(root=self.data_root,
train=train,
num_samples=num_samples,
download=True)
loader = DataLoader(dataset=dataset,
batch_size=self.batch_size,
num_workers=0,
shuffle=train)
return loader
def train_dataloader(self):
return DataLoader(TrialMNIST(root=_PATH_DATASETS,
train=True,
download=True),
batch_size=16)
def prepare_data(self):
TrialMNIST(root=self.data_root, train=True, download=True)
def prepare_data(self):
TrialMNIST(self.data_dir, train=True, download=True)
TrialMNIST(self.data_dir, train=False, download=True)
def train_dataloader(self):
return DataLoader(TrialMNIST(train=True, download=True), batch_size=16)