Data transforms now use v2

model_training/data_setup.py

@@ -1,7 +1,9 @@
 #####################################
 # Packages & Dependencies
 #####################################
-from torchvision import transforms, datasets
+import torch
+from torchvision import datasets
+from torchvision.transforms import v2
 from torch.utils.data import DataLoader
 
 import utils
@@ -14,18 +16,20 @@ import numpy as np
 import matplotlib.pyplot as plt
 
 # Transformations applied to each image
-BASE_TRANSFORMS = transforms.Compose([
-    transforms.ToTensor(), # Convert to tensor and rescale pixel values to within [0, 1]
-    transforms.Normalize(mean = [0.1307], std = [0.3081]) # Normalize with MNIST stats
+BASE_TRANSFORMS = v2.Compose([
+    v2.ToImage(), # Convert to tensor
+    v2.ToDtype(torch.float32, scale = True), # Rescale pixel values to within [0, 1]
+    v2.Normalize(mean = [0.1307], std = [0.3081]) # Normalize with MNIST stats
 ])
 
-TRAIN_TRANSFORMS = transforms.Compose([
-    transforms.RandomAffine(degrees = 15, # Rotate up to -/+ 15 degrees
-                            scale = (0.8, 1.2), # Scale between 80 and 120 percent
-                            translate = (0.08, 0.08), # Translate up to -/+ 8 percent in both x and y
-                            shear = 10),  # Shear up to -/+ 10 degrees
-    transforms.ToTensor(), # Convert to tensor and rescale pixel values to within [0, 1]
-    transforms.Normalize(mean = [0.1307], std = [0.3081]), # Normalize with MNIST stats
+TRAIN_TRANSFORMS = v2.Compose([
+    v2.RandomAffine(degrees = 15, # Rotate up to -/+ 15 degrees
+                    scale = (0.8, 1.2), # Scale between 80 and 120 percent
+                    translate = (0.08, 0.08), # Translate up to -/+ 8 percent in both x and y
+                    shear = 10),  # Shear up to -/+ 10 degrees
+    v2.ToImage(), # Convert to tensor
+    v2.ToDtype(torch.float32, scale = True), # Rescale pixel values to within [0, 1]
+    v2.Normalize(mean = [0.1307], std = [0.3081]), # Normalize with MNIST stats
 ])