From 1c708e9a449a90c866d7efb44798ba8a74ee4e85 Mon Sep 17 00:00:00 2001 From: Marcin Chrzanowski Date: Sat, 1 May 2021 15:01:02 +0200 Subject: Add helper to get both loaders --- src/data.py | 9 +++++++++ 1 file changed, 9 insertions(+) (limited to 'src') diff --git a/src/data.py b/src/data.py index ff7aca5..8019b88 100644 --- a/src/data.py +++ b/src/data.py @@ -141,3 +141,12 @@ def make_augmented_data(augmentations, after_target_transform=None): datasets.append(dataset) return torch.utils.data.ConcatDataset(datasets) + +def get_loaders(augmentations, target_transform, train_batch_size): + train_data = make_augmented_data(augmentations, after_target_transform=target_transform) + train_loader = torch.utils.data.DataLoader(train_data, batch_size=train_batch_size, shuffle=True) + + test_data = ShapesData(False, target_transform=target_transform) + test_loader = torch.utils.data.DataLoader(test_data, batch_size=1000, shuffle=False) + + return train_loader, test_loader -- cgit v1.2.3