Fix FID metric bug

Both the real and the fake images should be normalized in the [0, 1] range. Before this, our real images came normalized (with mean = std = 0.5).

Author: daniel-gallo

datasets/celeba.py

@@ -7,7 +7,7 @@
 from datasets.sampler import ResumableSeedableSampler
 
 
-def get_celeba_dataloader(batch_size, seed, data_dir="data/"):
+def get_celeba_dataloader(batch_size, seed, data_dir="data/", normalize: bool = True):
     """
     Builds a dataloader with all images from the CelebA dataset.
     Args:
@@ -18,17 +18,27 @@ def get_celeba_dataloader(batch_size, seed, data_dir="data/"):
         DataLoader: DataLoader object containing the dataset.
 
     """
-    mean = (0.5, 0.5, 0.5)
-    std = (0.5, 0.5, 0.5)
-
-    data_transforms = transforms.Compose(
-        [
-            transforms.ToTensor(),
-            transforms.Normalize(mean, std),
-            transforms.CenterCrop((178, 178)),
-            transforms.Resize((64, 64)),
-        ]
-    )
+
+    if normalize:
+        mean = (0.5, 0.5, 0.5)
+        std = (0.5, 0.5, 0.5)
+
+        data_transforms = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Normalize(mean, std),
+                transforms.CenterCrop((178, 178)),
+                transforms.Resize((64, 64)),
+            ]
+        )
+    else:
+        data_transforms = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.CenterCrop((178, 178)),
+                transforms.Resize((64, 64)),
+            ]
+        )
 
     path = Path(data_dir)
 

datasets/cifar10.py

@@ -7,11 +7,7 @@
 from datasets.sampler import ResumableSeedableSampler
 
 
-def get_cifar10_dataloader(
-    batch_size,
-    seed,
-    data_dir,
-):
+def get_cifar10_dataloader(batch_size, seed, data_dir, normalize: bool = True):
     """
     Builds a dataloader with all training images from the CIFAR-10 dataset.
     Args:
@@ -22,13 +18,15 @@ def get_cifar10_dataloader(
         DataLoader: DataLoader object containing the dataset.
 
     """
-
-    mean = (0.5, 0.5, 0.5)
-    std = (0.5, 0.5, 0.5)
-
-    transform = transforms.Compose(
-        [transforms.ToTensor(), transforms.Normalize(mean, std)]
-    )
+    if normalize:
+        mean = (0.5, 0.5, 0.5)
+        std = (0.5, 0.5, 0.5)
+
+        transform = transforms.Compose(
+            [transforms.ToTensor(), transforms.Normalize(mean, std)]
+        )
+    else:
+        transform = transforms.ToTensor()
 
     path = Path(data_dir) / "cifar10"
 

utils/evaluation_utils.py

@@ -1,16 +1,16 @@
 from pathlib import Path
 
-from PIL import Image
 import torch
+from PIL import Image
 from torchvision import transforms
 from torchvision.utils import save_image
 
-from datasets.cifar10 import get_cifar10_dataloader
 from datasets.celeba import get_celeba_dataloader
+from datasets.cifar10 import get_cifar10_dataloader
 
 
 def read_samples(path):
-    transform = transforms.Compose([transforms.Resize((32, 32)), transforms.ToTensor()])
+    transform = transforms.Compose([transforms.ToTensor()])
 
     tensor_list = []
     for p in Path(path).rglob("*.png"):
@@ -24,9 +24,9 @@ def read_samples(path):
 
 def get_dataset_samples(dataset_name, data_path, seed, n_samples):
     if dataset_name == "cifar10":
-        dataset = get_cifar10_dataloader(n_samples, seed, data_path)
+        dataset = get_cifar10_dataloader(n_samples, seed, data_path, normalize=False)
     elif dataset_name == "celeba":
-        dataset = get_celeba_dataloader(n_samples, seed, data_path)
+        dataset = get_celeba_dataloader(n_samples, seed, data_path, normalize=False)
     else:
         raise ValueError("Incorrect dataset name")