Integrate training uvit and deediff with imagenet64 and imagenet256

Author: razvanmatisan

datasets/imagenet.py

@@ -1,15 +1,18 @@
+from pathlib import Path
+
 from torch.utils.data import DataLoader
 from torchvision import datasets, transforms
 
-from pathlib import Path
 from datasets.sampler import ResumableSeedableSampler
 
 
 # https://www.kaggle.com/datasets/dimensi0n/imagenet-256
 def get_imagenet_dataloader(
     batch_size,
     seed,
-    data_dir="./archive",
+    data_dir,
+    resize: bool,  # resizing to 64x64
+    normalize: bool = True,
 ):
     """
     Builds a dataloader with all images from a 540k subset of ImageNet (with 256x256 resolution).
@@ -22,13 +25,19 @@ def get_imagenet_dataloader(
         DataLoader: DataLoader object containing the dataset.
     """
 
-    mean = (0.5, 0.5, 0.5)
-    std = (0.5, 0.5, 0.5)
-
     # All images from the dataset are 256x256 resolution
-    transform = transforms.Compose(
-        [transforms.ToTensor(), transforms.Normalize(mean, std)]
-    )
+    transformations = [transforms.ToTensor()]
+
+    if normalize:
+        mean = (0.5, 0.5, 0.5)
+        std = (0.5, 0.5, 0.5)
+
+        transformations.append(transforms.Normalize(mean, std))
+
+    if resize:
+        transformations.append(transforms.Resize((64, 64)))
+
+    transform = transforms.Compose(transformations)
 
     path = Path(data_dir) / "imagenet"
 

main.py

@@ -184,7 +184,7 @@ def get_args():
         "--dataset",
         type=str,
         default="cifar10",
-        choices=["cifar10", "celeba", "imagenet"],
+        choices=["cifar10", "celeba", "imagenet64", "imagenet256"],
         help="Dataset name",
     )
     parser.add_argument(
@@ -204,7 +204,7 @@ def main():
         config = load_config(args.config_path)
         args.__dict__.update(config["model_params"])
 
-    if args.dataset == "imagenet":
+    if args.dataset == "imagenet256":
         args.__dict__.update(config["autoencoder"])
 
     torch.use_deterministic_algorithms(True)

tests/test_datasets.py

@@ -22,22 +22,12 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-# Might delete later
-def ignore_if_imagenet_data_not_downloaded(f):
-    @wraps(f)
-    def wrapper(*args, **kwargs):
-        if not Path("archive/").exists():
-            return
-
-        return f(*args, **kwargs)
-
-    return wrapper
-
-
 @ignore_if_data_not_downloaded
 @pytest.mark.parametrize("batch_size", [16])
 def test_cifar10(batch_size):
-    dataloader = get_cifar10_dataloader(batch_size=batch_size, seed=0)
+    dataloader = get_cifar10_dataloader(
+        batch_size=batch_size, seed=0, data_dir="./data"
+    )
 
     x, _ = next(iter(dataloader))
     assert x.shape == torch.Size([batch_size, 3, 32, 32])
@@ -46,16 +36,29 @@ def test_cifar10(batch_size):
 @ignore_if_data_not_downloaded
 @pytest.mark.parametrize("batch_size", [4])
 def test_celeba(batch_size):
-    dataloader = get_celeba_dataloader(batch_size=batch_size, seed=0)
+    dataloader = get_celeba_dataloader(batch_size=batch_size, seed=0, data_dir="./data")
 
     x, _ = next(iter(dataloader))
     assert x.shape == torch.Size([batch_size, 3, 64, 64])
 
 
-@ignore_if_imagenet_data_not_downloaded
+@ignore_if_data_not_downloaded
+@pytest.mark.parametrize("batch_size", [4])
+def test_imagenet64(batch_size):
+    dataloader = get_imagenet_dataloader(
+        batch_size=batch_size, seed=0, data_dir="./data", resize=True
+    )
+
+    x, _ = next(iter(dataloader))
+    assert x.shape == torch.Size([batch_size, 3, 64, 64])
+
+
+@ignore_if_data_not_downloaded
 @pytest.mark.parametrize("batch_size", [4])
-def test_imagenet(batch_size):
-    dataloader = get_imagenet_dataloader(batch_size=batch_size, seed=0)
+def test_imagenet256(batch_size):
+    dataloader = get_imagenet_dataloader(
+        batch_size=batch_size, seed=0, data_dir="./data", resize=False
+    )
 
     x, _ = next(iter(dataloader))
     assert x.shape == torch.Size([batch_size, 3, 256, 256])

trainer.py

@@ -140,11 +140,19 @@ def _init_dataloader(self):
                 seed=self.args.seed,
                 data_dir=self.args.data_path,
             )
-        elif self.args.dataset == "imagenet":
+        elif self.args.dataset == "imagenet64":
             self.dataloader = get_imagenet_dataloader(
                 batch_size=self.args.batch_size,
                 seed=self.args.seed,
                 data_dir=self.args.data_path,
+                resize=True,
+            )
+        elif self.args.dataset == "imagenet256":
+            self.dataloader = get_imagenet_dataloader(
+                batch_size=self.args.batch_size,
+                seed=self.args.seed,
+                data_dir=self.args.data_path,
+                resize=False,
             )
         else:
             raise ValueError(f"Dataset {self.args.dataset} not implemented.")
@@ -298,7 +306,7 @@ def _loss_fn(self, batch):
         data = batch[0].to(self.device)
         batch_size = data.size(0)
         clean_images = data
-        labels = batch[1].to(self.device) if self.args.dataset == "imagenet" else None
+        labels = batch[1].to(self.device) if "imagenet" in self.args.dataset else None
 
         timesteps = torch.randint(
             0, self.args.num_timesteps, (batch_size,), device=self.device
@@ -308,13 +316,13 @@ def _loss_fn(self, batch):
 
         if self.args.model == "uvit":
             if self.args.parametrization == "predict_noise":
-                predicted_noise = self.model(noisy_images, timesteps)
+                predicted_noise = self.model(noisy_images, timesteps, labels)
                 loss = F.mse_loss(predicted_noise, noise)
             elif self.args.parametrization == "predict_original":
-                predicted_original = self.model(noisy_images, timesteps)
+                predicted_original = self.model(noisy_images, timesteps, labels)
                 loss = F.mse_loss(predicted_original, clean_images)
             elif self.args.parametrization == "predict_previous":
-                predicted_previous = self.model(noisy_images, timesteps)
+                predicted_previous = self.model(noisy_images, timesteps, labels)
 
                 betas = torch.linspace(1e-4, 0.02, 1000).to(self.device)
                 alphas = 1 - betas