Add rotated MNIST dataset

.gitignore

@@ -1,4 +1,5 @@
 __pycache__/
 .vscode/
 *.pth
-scripts/
+scripts/
+data/

lib/datasets/__init__.py

@@ -3,6 +3,7 @@
 from torch.utils.data import random_split, Dataset
 
 from .synthetic import SyntheticDataset
+from .mnist import MNIST
 
 
 def train_val_split(dataset: Dataset, val_percentage: float) -> Tuple[Dataset]:

lib/datasets/mnist.py

@@ -0,0 +1,20 @@
+from torchvision import transforms
+from torchvision.datasets import mnist
+
+
+class MNIST(mnist.MNIST):
+    def __init__(
+        self,
+        rotated: bool = False,
+        root: str = "./data",
+        download: bool = True,
+        *args,
+        **kwargs
+    ):
+        transform = [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
+        if rotated:
+            transform.append(transforms.RandomRotation(degrees=45))
+        transform = transforms.Compose(transform)
+        super().__init__(
+            root=root, transform=transform, download=download, *args, **kwargs
+        )

requirements.txt

@@ -4,4 +4,5 @@ tqdm==4.63.0
 matplotlib==3.5.1
 seaborn==0.11.2
 -f https://download.pytorch.org/whl/torch_stable.html
-torch==1.9.0+cu102
+torch==1.9.0+cu102
+torchvision==0.10.0+cu102