Added Fashion MNIST support

ancestor-mithril · May 22, 2024 · 1b4e77f · 1b4e77f
1 parent 11c4f56
commit 1b4e77f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/utils/dataset.py b/utils/dataset.py
@@ -1,7 +1,7 @@
 from functools import partial
 
 from torch.utils.data import Dataset, DataLoader
-from torchvision.datasets import CIFAR10, CIFAR100
+from torchvision.datasets import CIFAR10, CIFAR100, FashionMNIST
 
 from .transforms import init_transforms
 
@@ -36,6 +36,9 @@ def init_dataset(args):
     elif args.dataset == 'cifar100':
         dataset_fn = partial(CIFAR100, root=args.data_path, download=True)
         num_classes = 100
+    elif args.dataset == 'FashionMNIST':
+        dataset_fn = partial(FashionMNIST, root=args.data_path, download=True)
+        num_classes = 10
     else:
         raise NotImplementedError(f'Dataset {args.dataset} not implemented')
 

diff --git a/utils/transforms.py b/utils/transforms.py
@@ -58,6 +58,6 @@ def test_runtime(self):
 
 
 def init_transforms(args) -> DatasetTransforms:
-    if args.dataset in ('cifar10', 'cifar100'):
+    if args.dataset in ('cifar10', 'cifar100', 'FashionMNIST'):
         return CifarTransforms(args)
     raise NotImplementedError(f"Transforms not implemented for {args.dataset}")