separated train and test transformation

Data.py

@@ -2,12 +2,18 @@ from torchvision import datasets, transforms, models
 import torch
 import numpy as np
 
-# transform images to size
-def transform(res):
+# train set transform 
+def train_transform(res):
     return transforms.Compose([
         # ResNet expects 224 x 224 res 
+        # Inception expects 299 x 299
         transforms.Resize((res, res)),
-        transforms.RandomHorizontalFlip(),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.ColorJitter(
+            brightness=0.2, 
+            contrast=0.2, 
+            saturation=0.1
+        ),
         transforms.ToTensor(),
         # normalise to 
         transforms.Normalize(
@@ -16,14 +22,26 @@ def transform(res):
         )
     ])
 
+# test set transform
+def test_transform(res):
+    return transforms.Compose([
+        # Just standard resize to 224x224
+        transforms.Resize((res, res)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], 
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
 # Load data with 'identity' as target and transform it
-def get_set(res):
+def get_set():
     return datasets.CelebA(
         root='./data', 
         split='all', 
         target_type='identity', 
         download=True, 
-        transform=transform(res)
+        transform=None
     )
 
 
@@ -66,7 +84,6 @@ def get_indices(dataset, identities, split_at):
     test_indices = []
 
     #training_sample = int(sample_size * training_ratio)
-
     for person_id in identities:
         # Get all indices for this specific person
         indices = torch.where(dataset.identity == person_id)[0].numpy()