IJEPA — Image Joint-Embedding Predictive Architecture

Minimal PyTorch implementation of I-JEPA (ICLR 2023, Y. LeCun et al.). Self-supervised visual representation learning via multi-block masking + representation prediction.

Setup

pip install torch torchvision pillow numpy

No timm, no Hydra, no Lightning, no YAML.

Usage

Train

Place images anywhere under data/train/:

data/train/
  any_folder/
    img001.jpg
    img002.jpg
    ...

python train.py

Checkpoints saved to checkpoints/model.pt.

Inference

python inference.py --checkpoint checkpoints/model.pt --image dog.jpg

Prints CLS shape: (192,) and Patch shape: (196, 192). Optional --save_npy to dump patch embeddings.

Architecture

Component	Encoder	Predictor
Depth	12	6
Dim	192	192
Heads	3	3
Params	~5.7M	~2.9M

1. Masking — Sample 1 large context block (85–100%) + 4 small target blocks (15–20%) per image. Overlapping patches removed from context.
2. Context encoder — ViT on visible context patches only.
3. Predictor — Narrow ViT, called once per target block. Receives context embeddings + mask tokens (shared learnable vector + positional embedding). Predicts target block representations.
4. Loss — 1/M Σᵢ MSE(pred_blockᵢ, target_blockᵢ)
5. Target encoder — EMA copy of context encoder (momentum 0.996).

Config

Edit config.py:

EMBED_DIM = 192   # 192 / 384 / 768
DEPTH = 12         # ViT depth
BATCH_SIZE = 64
EPOCHS = 100

Acknowledgments

• Original paper: I-JEPA by Assran et al.