change/add files for ijepa training

mmlearn/modules/encoders/vision.py

@@ -561,6 +561,10 @@ def forward(
         return self.predictor_proj(x)
 
 
+@store(
+    group="modules/encoders",
+    provider="mmlearn",
+)
 def vit_predictor(**kwargs: Any) -> VisionTransformerPredictor:
     """
     Create a VisionTransformerPredictor model.
@@ -575,6 +579,10 @@ def vit_predictor(**kwargs: Any) -> VisionTransformerPredictor:
     )
 
 
+@store(
+    group="modules/encoders",
+    provider="mmlearn",
+)
 def vit_tiny(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
     """
     Create a VisionTransformer model with tiny configuration.

mmlearn/tasks/ijepa_pretraining.py

@@ -5,6 +5,7 @@
 import lightning as L  # noqa: N812
 import torch
 import torch.nn.functional as F  # noqa: N812
+from hydra_zen import store
 from lightning.pytorch.utilities.types import OptimizerLRScheduler
 from lightning_utilities.core.rank_zero import rank_zero_warn
 
@@ -15,6 +16,7 @@
 from mmlearn.modules.encoders.vision import VisionTransformer
 
 
+@store(group="task", provider="mmlearn")
 class IJEPA(L.LightningModule):
     """Pretraining module for IJEPA.
 

projects/ijepa/configs/experiment/run_ijepa.yaml

@@ -0,0 +1,74 @@
+# @package _global_
+
+defaults:
+  - /[email protected]: ImageNet
+  - /datasets/[email protected]: med_clip_vision_transform
+  - /[email protected]: ImageNet
+  - /datasets/[email protected]: med_clip_vision_transform
+  - /modules/[email protected]: vit_tiny
+  - /modules/[email protected]: vit_predictor
+  - /modules/[email protected]: AdamW
+  - /modules/[email protected]_scheduler.scheduler: CosineAnnealingLR
+  - /trainer/[email protected]_monitor: LearningRateMonitor
+  - /trainer/[email protected]_checkpoint: ModelCheckpoint
+  - /trainer/[email protected]_stopping: EarlyStopping
+  - /trainer/[email protected]_summary: ModelSummary
+  - /trainer/[email protected]: WandbLogger
+  - override /task: IJEPA
+  - _self_
+
+seed: 0
+
+datasets:
+  val:
+    split: valid
+    transform:
+      job_type: eval
+
+dataloader:
+  train:
+    batch_size: 32
+    num_workers: 4
+  val:
+    batch_size: 32
+    num_workers: 4
+
+task:
+  optimizer:
+    betas:
+    - 0.9
+    - 0.98
+    lr: 5.0e-5
+    weight_decay: 0.1
+    eps: 1.0e-6
+  lr_scheduler:
+    scheduler:
+      T_max: 107_559 # make sure to change this if max_epochs or accumulate_grad_batches is changed
+    extras:
+      interval: step
+
+trainer:
+  max_epochs: 20
+  precision: 16-mixed
+  deterministic: False
+  benchmark: True
+  sync_batchnorm: False # set to True if using DDP with batchnorm
+  log_every_n_steps: 100
+  accumulate_grad_batches: 4
+  check_val_every_n_epoch: 1
+  callbacks:
+    model_checkpoint:
+      monitor: val/loss
+      save_top_k: 1
+      save_last: True
+      every_n_epochs: 1
+    early_stopping:
+      monitor: val/loss
+      patience: 5
+      mode: min
+    model_summary:
+      max_depth: 2
+
+tags:
+  - ${experiment_name}
+  - ijepa pretraining