Added IJEPA model

Author: ParamThakkar123

jepa_try.py

@@ -0,0 +1,7 @@
+from world_models.models.jepa_agent import JEPAAgent
+
+agent = JEPAAgent(
+    folder="results/jepa_try",
+    write_tag="jepa_try",
+)
+agent.train()

world_models/configs/jepa_config.py

@@ -0,0 +1,108 @@
+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.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)
+
+        # 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": {
+                "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,
+            },
+            "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/imagenet1k.py

@@ -0,0 +1,214 @@
+import os
+import subprocess
+import time
+
+import numpy as np
+
+from logging import getLogger
+
+import torch
+import torchvision
+
+_GLOBAL_SEED = 0
+logger = getLogger()
+
+
+def make_imagenet1k(
+    transform,
+    batch_size,
+    collator=None,
+    pin_mem=True,
+    num_workers=8,
+    world_size=1,
+    rank=0,
+    root_path=None,
+    image_folder=None,
+    training=True,
+    copy_data=False,
+    drop_last=True,
+    subset_file=None,
+):
+    dataset = ImageNet(
+        root=root_path,
+        image_folder=image_folder,
+        transform=transform,
+        train=training,
+        copy_data=copy_data,
+        index_targets=False,
+    )
+    if subset_file is not None:
+        dataset = ImageNetSubset(dataset, subset_file)
+    logger.info("ImageNet dataset created")
+    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("ImageNet unsupervised data loader created")
+
+    return dataset, data_loader, dist_sampler
+
+
+class ImageNet(torchvision.datasets.ImageFolder):
+
+    def __init__(
+        self,
+        root,
+        image_folder="imagenet_full_size/061417/",
+        tar_file="imagenet_full_size-061417.tar.gz",
+        transform=None,
+        train=True,
+        job_id=None,
+        local_rank=None,
+        copy_data=True,
+        index_targets=False,
+    ):
+        """
+        ImageNet
+
+        Dataset wrapper (can copy data locally to machine)
+
+        :param root: root network directory for ImageNet data
+        :param image_folder: path to images inside root network directory
+        :param tar_file: zipped image_folder inside root network directory
+        :param train: whether to load train data (or validation)
+        :param job_id: scheduler job-id used to create dir on local machine
+        :param copy_data: whether to copy data from network file locally
+        :param index_targets: whether to index the id of each labeled image
+        """
+
+        suffix = "train/" if train else "val/"
+        data_path = None
+        if copy_data:
+            logger.info("copying data locally")
+            data_path = copy_imgnt_locally(
+                root=root,
+                suffix=suffix,
+                image_folder=image_folder,
+                tar_file=tar_file,
+                job_id=job_id,
+                local_rank=local_rank,
+            )
+        if (not copy_data) or (data_path is None):
+            data_path = os.path.join(root, image_folder, suffix)
+        logger.info(f"data-path {data_path}")
+
+        super(ImageNet, self).__init__(root=data_path, transform=transform)
+        logger.info("Initialized ImageNet")
+
+        if index_targets:
+            self.targets = []
+            for sample in self.samples:
+                self.targets.append(sample[1])
+            self.targets = np.array(self.targets)
+            self.samples = np.array(self.samples)
+
+            mint = None
+            self.target_indices = []
+            for t in range(len(self.classes)):
+                indices = np.squeeze(np.argwhere(self.targets == t)).tolist()
+                self.target_indices.append(indices)
+                mint = len(indices) if mint is None else min(mint, len(indices))
+                logger.debug(f"num-labeled target {t} {len(indices)}")
+            logger.info(f"min. labeled indices {mint}")
+
+
+class ImageNetSubset(object):
+
+    def __init__(self, dataset, subset_file):
+        """
+        ImageNetSubset
+
+        :param dataset: ImageNet dataset object
+        :param subset_file: '.txt' file containing IDs of IN1K images to keep
+        """
+        self.dataset = dataset
+        self.subset_file = subset_file
+        self.filter_dataset_(subset_file)
+
+    def filter_dataset_(self, subset_file):
+        """Filter self.dataset to a subset"""
+        root = self.dataset.root
+        class_to_idx = self.dataset.class_to_idx
+        # -- update samples to subset of IN1k targets/samples
+        new_samples = []
+        logger.info(f"Using {subset_file}")
+        with open(subset_file, "r") as rfile:
+            for line in rfile:
+                class_name = line.split("_")[0]
+                target = class_to_idx[class_name]
+                img = line.split("\n")[0]
+                new_samples.append((os.path.join(root, class_name, img), target))
+        self.samples = new_samples
+
+    @property
+    def classes(self):
+        return self.dataset.classes
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        path, target = self.samples[index]
+        img = self.dataset.loader(path)
+        if self.dataset.transform is not None:
+            img = self.dataset.transform(img)
+        if self.dataset.target_transform is not None:
+            target = self.dataset.target_transform(target)
+        return img, target
+
+
+def copy_imgnt_locally(
+    root,
+    suffix,
+    image_folder="imagenet_full_size/061417/",
+    tar_file="imagenet_full_size-061417.tar.gz",
+    job_id=None,
+    local_rank=None,
+):
+    if job_id is None:
+        try:
+            job_id = os.environ["SLURM_JOBID"]
+        except Exception:
+            logger.info("No job-id, will load directly from network file")
+            return None
+
+    if local_rank is None:
+        try:
+            local_rank = int(os.environ["SLURM_LOCALID"])
+        except Exception:
+            logger.info("No job-id, will load directly from network file")
+            return None
+
+    source_file = os.path.join(root, tar_file)
+    target = f"/scratch/slurm_tmpdir/{job_id}/"
+    target_file = os.path.join(target, tar_file)
+    data_path = os.path.join(target, image_folder, suffix)
+    logger.info(f"{source_file}\n{target}\n{target_file}\n{data_path}")
+
+    tmp_sgnl_file = os.path.join(target, "copy_signal.txt")
+
+    if not os.path.exists(data_path):
+        if local_rank == 0:
+            commands = [["tar", "-xf", source_file, "-C", target]]
+            for cmnd in commands:
+                start_time = time.time()
+                logger.info(f"Executing {cmnd}")
+                subprocess.run(cmnd)
+                logger.info(f"Cmnd took {(time.time()-start_time)/60.} min.")
+            with open(tmp_sgnl_file, "+w") as f:
+                print("Done copying locally.", file=f)
+        else:
+            while not os.path.exists(tmp_sgnl_file):
+                time.sleep(60)
+                logger.info(f"{local_rank}: Checking {tmp_sgnl_file}")
+
+    return data_path