Changes and add datasets

Author: ParamThakkar123

.gitignore

@@ -4,3 +4,4 @@ results/
 dist/*
 pytorch_world.egg-info/
 world_models/models/data/
+cifar

jepa_try.py

@@ -1,7 +1,21 @@
 from world_models.models.jepa_agent import JEPAAgent
 
-agent = JEPAAgent(
-    folder="results/jepa_try",
-    write_tag="jepa_try",
-)
-agent.train()
+if __name__ == "__main__":
+    agent = JEPAAgent(
+        dataset="cifar10",
+        root_path=r"E:\pytorch-world\cifar",
+        download=True,
+        folder="results/cifar_jepa",
+        write_tag="cifar_jepa",
+        batch_size=16,
+        pin_mem=False,
+        crop_size=32,
+        patch_size=4,
+        enc_mask_scale=(0.05, 0.15),
+        pred_mask_scale=(0.05, 0.15),
+        min_keep=1,
+        allow_overlap=True,
+        num_workers=0,
+        epochs=25,
+    )
+    agent.train()

world_models/configs/jepa_config.py

@@ -19,6 +19,10 @@ def __init__(self):
         self.pred_emb_dim: int = 384
 
         # data
+        self.dataset: str = "imagenet"  # "imagenet" or "imagefolder"
+        self.val_split: float | None = (
+            None  # optional fraction for val split when using imagefolder
+        )
         self.use_gaussian_blur: bool = True
         self.use_horizontal_flip: bool = True
         self.use_color_distortion: bool = True
@@ -30,6 +34,7 @@ def __init__(self):
         self.image_folder: str = "train"
         self.crop_size: int = 224
         self.crop_scale: Tuple[float, float] = (0.67, 1.0)
+        self.download: bool = False  # allow CIFAR10 download if missing
 
         # mask
         self.allow_overlap: bool = False
@@ -68,6 +73,8 @@ def to_dict(self) -> Dict[str, Dict[str, Any]]:
                 "pred_emb_dim": self.pred_emb_dim,
             },
             "data": {
+                "dataset": self.dataset,
+                "val_split": self.val_split,
                 "use_gaussian_blur": self.use_gaussian_blur,
                 "use_horizontal_flip": self.use_horizontal_flip,
                 "use_color_distortion": self.use_color_distortion,
@@ -79,6 +86,7 @@ def to_dict(self) -> Dict[str, Dict[str, Any]]:
                 "image_folder": self.image_folder,
                 "crop_size": self.crop_size,
                 "crop_scale": self.crop_scale,
+                "download": self.download,  # new
             },
             "mask": {
                 "allow_overlap": self.allow_overlap,

world_models/datasets/cifar10.py

@@ -0,0 +1,41 @@
+import torch
+from torchvision.datasets import CIFAR10
+from logging import getLogger
+
+logger = getLogger()
+
+
+def make_cifar10(
+    transform,
+    batch_size,
+    collator=None,
+    pin_mem=True,
+    num_workers=8,
+    world_size=1,
+    rank=0,
+    root_path=None,
+    drop_last=True,
+    train=True,
+    download=False,  # new
+):
+    dataset = CIFAR10(
+        root=root_path,
+        train=train,
+        download=download,
+        transform=transform,
+    )
+    dist_sampler = torch.utils.data.distributed.DistributedSampler(
+        dataset=dataset, num_replicas=world_size, rank=rank
+    )
+    data_loader = torch.utils.data.DataLoader(
+        dataset,
+        collate_fn=collator,
+        sampler=dist_sampler,
+        batch_size=batch_size,
+        drop_last=drop_last,
+        pin_memory=pin_mem,
+        num_workers=num_workers,
+        persistent_workers=False,
+    )
+    logger.info("CIFAR10 data loader created")
+    return dataset, data_loader, dist_sampler

world_models/datasets/imagenet1k.py

@@ -8,6 +8,7 @@
 
 import torch
 import torchvision
+from torch.utils.data import random_split
 
 _GLOBAL_SEED = 0
 logger = getLogger()
@@ -212,3 +213,41 @@ def copy_imgnt_locally(
                 logger.info(f"{local_rank}: Checking {tmp_sgnl_file}")
 
     return data_path
+
+
+def make_imagefolder(
+    transform,
+    batch_size,
+    collator=None,
+    pin_mem=True,
+    num_workers=8,
+    world_size=1,
+    rank=0,
+    root_path=None,
+    image_folder=None,
+    drop_last=True,
+    val_split: float | None = None,
+):
+    dataset = torchvision.datasets.ImageFolder(
+        root=os.path.join(root_path, image_folder) if image_folder else root_path,
+        transform=transform,
+    )
+    if val_split:
+        val_size = int(len(dataset) * val_split)
+        train_size = len(dataset) - val_size
+        dataset, _ = random_split(dataset, [train_size, val_size])
+    dist_sampler = torch.utils.data.distributed.DistributedSampler(
+        dataset=dataset, num_replicas=world_size, rank=rank
+    )
+    data_loader = torch.utils.data.DataLoader(
+        dataset,
+        collate_fn=collator,
+        sampler=dist_sampler,
+        batch_size=batch_size,
+        drop_last=drop_last,
+        pin_memory=pin_mem,
+        num_workers=num_workers,
+        persistent_workers=False,
+    )
+    logger.info("ImageFolder data loader created")
+    return dataset, data_loader, dist_sampler

world_models/training/train_jepa.py

@@ -27,7 +27,8 @@
     AverageMeter,
 )
 from world_models.utils.jepa_utils import repeat_interleave_batch
-from world_models.datasets.imagenet1k import make_imagenet1k
+from world_models.datasets.imagenet1k import make_imagenet1k, make_imagefolder
+from world_models.datasets.cifar10 import make_cifar10
 from world_models.helpers.jepa_helper import load_checkpoint, init_model, init_opt
 from world_models.transforms.transforms import make_transforms
 from world_models.configs.jepa_config import JEPAConfig
@@ -181,20 +182,52 @@ def main(args, resume_preempt=False):
     )
 
     # -- init data-loaders/samplers
-    _, unsupervised_loader, unsupervised_sampler = make_imagenet1k(
-        transform=transform,
-        batch_size=batch_size,
-        collator=mask_collator,
-        pin_mem=pin_mem,
-        training=True,
-        num_workers=num_workers,
-        world_size=world_size,
-        rank=rank,
-        root_path=root_path,
-        image_folder=image_folder,
-        copy_data=copy_data,
-        drop_last=True,
-    )
+    dataset_type = args["data"]["dataset"]
+    val_split = args["data"]["val_split"]
+    download = args["data"].get("download", False)
+    if dataset_type.lower() == "imagenet":
+        _, unsupervised_loader, unsupervised_sampler = make_imagenet1k(
+            transform=transform,
+            batch_size=batch_size,
+            collator=mask_collator,
+            pin_mem=pin_mem,
+            training=True,
+            num_workers=num_workers,
+            world_size=world_size,
+            rank=rank,
+            root_path=root_path,
+            image_folder=image_folder,
+            copy_data=copy_data,
+            drop_last=True,
+        )
+    elif dataset_type.lower() == "cifar10":
+        _, unsupervised_loader, unsupervised_sampler = make_cifar10(
+            transform=transform,
+            batch_size=batch_size,
+            collator=mask_collator,
+            pin_mem=pin_mem,
+            num_workers=num_workers,
+            world_size=world_size,
+            rank=rank,
+            root_path=root_path,
+            drop_last=True,
+            train=True,
+            download=download,  # pass through
+        )
+    else:
+        _, unsupervised_loader, unsupervised_sampler = make_imagefolder(
+            transform=transform,
+            batch_size=batch_size,
+            collator=mask_collator,
+            pin_mem=pin_mem,
+            num_workers=num_workers,
+            world_size=world_size,
+            rank=rank,
+            root_path=root_path,
+            image_folder=image_folder,
+            drop_last=True,
+            val_split=val_split,
+        )
     ipe = len(unsupervised_loader)
 
     # -- init optimizer and scheduler
@@ -212,9 +245,17 @@ def main(args, resume_preempt=False):
         ipe_scale=ipe_scale,
         use_bfloat16=use_bfloat16,
     )
-    encoder = DistributedDataParallel(encoder, static_graph=True)
-    predictor = DistributedDataParallel(predictor, static_graph=True)
-    target_encoder = DistributedDataParallel(target_encoder)
+
+    is_distributed = (
+        torch.distributed.is_available()
+        and torch.distributed.is_initialized()
+        and world_size > 1
+    )
+    if is_distributed:
+        encoder = DistributedDataParallel(encoder, static_graph=True)
+        predictor = DistributedDataParallel(predictor, static_graph=True)
+        target_encoder = DistributedDataParallel(target_encoder)
+    # keep modules unwrapped when not distributed
     for p in target_encoder.parameters():
         p.requires_grad = False
 
@@ -328,7 +369,8 @@ def loss_fn(z, h):
                 else:
                     loss.backward()
                     optimizer.step()
-                grad_stats = grad_logger(encoder.named_parameters())
+                enc_for_log = encoder.module if is_distributed else encoder
+                grad_stats = grad_logger(enc_for_log.named_parameters())
                 optimizer.zero_grad()
 
                 # Step 3. momentum update of target encoder

world_models/transforms/transforms.py

@@ -53,5 +53,7 @@ def __call__(self, img):
         if torch.bernoulli(torch.tensor(self.prob)) == 0:
             return img
 
-        radius = self.radius_min + torch.rand(1) * (self.radius_max - self.radius_min)
+        radius = self.radius_min + torch.rand(1).item() * (
+            self.radius_max - self.radius_min
+        )
         return img.filter(ImageFilter.GaussianBlur(radius=radius))

world_models/utils/utils.py

@@ -695,6 +695,6 @@ def max_episode_steps(self):
 def apply_masks(x, masks):
     all_x = []
     for m in masks:
-        mask_keep = m.unsqueeze(-1).repeat(1, 1, x.shape(-1))
-        all_x += [torch.gather(x, 1, mask_keep)]
+        mask_keep = m.unsqueeze(-1).repeat(1, 1, x.shape[-1])
+        all_x.append(torch.gather(x, 1, mask_keep))
     return torch.cat(all_x, dim=0)