Fid for iamgenet

Author: razvanmatisan

fid.py

@@ -12,7 +12,7 @@ def get_args():
         "--dataset",
         type=str,
         required=True,
-        choices=["cifar10", "celeba"],
+        choices=["cifar10", "celeba", "imagenet64", "imagenet256"],
         help="Dataset name.",
     )
     parser.add_argument(
@@ -43,6 +43,7 @@ def main():
     args = get_args()
     generated_images = read_samples(args.samples_path)
     n_samples = len(generated_images)
+    print(f"Using {n_samples}")
     real_images = get_dataset_samples(
         args.dataset, args.data_path, args.seed, n_samples
     )

utils/evaluation_utils.py

@@ -7,6 +7,7 @@
 
 from datasets.celeba import get_celeba_dataloader
 from datasets.cifar10 import get_cifar10_dataloader
+from datasets.imagenet import get_imagenet_dataloader
 
 
 def read_samples(path):
@@ -28,6 +29,10 @@ def get_dataset_samples(dataset_name, data_path, seed, n_samples):
         dataset = get_cifar10_dataloader(n_samples, seed, data_path, normalize=False)
     elif dataset_name == "celeba":
         dataset = get_celeba_dataloader(n_samples, seed, data_path, normalize=False)
+    elif dataset_name == "imagenet64":
+        dataset = get_imagenet_dataloader(n_samples, seed, data_path, normalize=False, resize=True)
+    elif dataset_name == "imagenet256":
+        dataset = get_imagenet_dataloader(n_samples, seed, data_path, normalize=False, resize=False)
     else:
         raise ValueError("Incorrect dataset name")