Merge pull request #8 from ParamThakkar123/add_jepa

Added IJEPA model

Author: ParamThakkar123

.gitignore

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

jepa_try.py

@@ -0,0 +1,21 @@
+from world_models.models.jepa_agent import JEPAAgent
+
+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

@@ -0,0 +1,116 @@
+import os
+from typing import Tuple, Dict, Any
+
+
+class JEPAConfig:
+    """
+    Minimal configuration container for JEPA training.
+    Converts to the nested dict expected by `train_jepa.main`.
+    """
+
+    def __init__(self):
+        # meta
+        self.use_bfloat16: bool = False
+        self.model_name: str = "vit_base"
+        self.load_checkpoint: bool = False
+        self.read_checkpoint: str | None = None
+        self.copy_data: bool = False
+        self.pred_depth: int = 6
+        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
+        self.color_jitter_strength: float = 0.5
+        self.batch_size: int = 64
+        self.pin_mem: bool = True
+        self.num_workers: int = 8
+        self.root_path: str = os.environ.get("IMAGENET_ROOT", "/data/imagenet")
+        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
+        self.patch_size: int = 16
+        self.num_enc_masks: int = 1
+        self.min_keep: int = 4
+        self.enc_mask_scale: Tuple[float, float] = (0.15, 0.2)
+        self.num_pred_masks: int = 1
+        self.pred_mask_scale: Tuple[float, float] = (0.15, 0.2)
+        self.aspect_ratio: Tuple[float, float] = (0.75, 1.5)
+
+        # optimization
+        self.ema: Tuple[float, float] = (0.996, 1.0)
+        self.ipe_scale: float = 1.0
+        self.weight_decay: float = 0.04
+        self.final_weight_decay: float = 0.4
+        self.epochs: int = 300
+        self.warmup: int = 40
+        self.start_lr: float = 1e-6
+        self.lr: float = 1.5e-4
+        self.final_lr: float = 1e-6
+
+        # logging
+        self.folder: str = "results/jepa"
+        self.write_tag: str = "jepa_run"
+
+    def to_dict(self) -> Dict[str, Dict[str, Any]]:
+        return {
+            "meta": {
+                "use_bfloat16": self.use_bfloat16,
+                "model_name": self.model_name,
+                "load_checkpoint": self.load_checkpoint,
+                "read_checkpoint": self.read_checkpoint,
+                "copy_data": self.copy_data,
+                "pred_depth": self.pred_depth,
+                "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,
+                "color_jitter_strength": self.color_jitter_strength,
+                "batch_size": self.batch_size,
+                "pin_mem": self.pin_mem,
+                "num_workers": self.num_workers,
+                "root_path": self.root_path,
+                "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,
+                "patch_size": self.patch_size,
+                "num_enc_masks": self.num_enc_masks,
+                "min_keep": self.min_keep,
+                "enc_mask_scale": self.enc_mask_scale,
+                "num_pred_masks": self.num_pred_masks,
+                "pred_mask_scale": self.pred_mask_scale,
+                "aspect_ratio": self.aspect_ratio,
+            },
+            "optimization": {
+                "ema": self.ema,
+                "ipe_scale": self.ipe_scale,
+                "weight_decay": self.weight_decay,
+                "final_weight_decay": self.final_weight_decay,
+                "epochs": self.epochs,
+                "warmup": self.warmup,
+                "start_lr": self.start_lr,
+                "lr": self.lr,
+                "final_lr": self.final_lr,
+            },
+            "logging": {
+                "folder": self.folder,
+                "write_tag": self.write_tag,
+            },
+        }

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