[WIP] Feature/rae

.gitignore

@@ -27,8 +27,6 @@ share/python-wheels/
 MANIFEST
 
 # PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
 *.manifest
 *.spec
 
@@ -83,36 +81,12 @@ notebooks/
 profile_default/
 ipython_config.py
 
-# pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-
 # pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#pdm.lock
-#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-#   in version control.
-#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
 .pdm.toml
 .pdm-python
 .pdm-build/
 
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+# PEP 582
 __pypackages__/
 
 # Celery stuff
@@ -155,13 +129,6 @@ dmypy.json
 # Cython debug symbols
 cython_debug/
 
-# PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
-
 # running logs
 examples/wandb
 outputs/

README.md

@@ -63,10 +63,10 @@ Here is a To-Do list, feel welcome to help to any point along this list. The alr
 - [ ] add some more context embedders
 - [ ] add reflow algorithm
 - [ ] add EDM
-- [ ] think about how to add a sampler abstraction and use it in the different Diffusion classes (generalist class with euler, heuns etc)
+- [x] think about how to add a sampler abstraction and use it in the different Diffusion classes (generalist class with euler, heuns etc)
 - [ ] Train our models on toy datasets for different tasks (conditional generation, Image to Image ...)
 - [ ] Add possibility to train LORA/DORA
-- [ ] add different sampler
+- [x] add different sampler
 - [ ] Try out Differential Transformers
 - [ ] Check to add https://arxiv.org/pdf/2406.02507
 - [ ] inject lessons learned from nvidia https://developer.nvidia.com/blog/rethinking-how-to-train-diffusion-models/

configs/dataset/imagenet_repa.yaml → configs/dataset/imagenet_latents.yaml

@@ -1,13 +1,13 @@
 # CIFAR10 dataset configuration
 train:
-  _target_: diffulab.datasets.ImageNetLatentREPA
+  _target_: diffulab.datasets.ImageNetLatent
   data_path: "data/imagenet"
   local: true
   batch_size: 128
   split: "train"
 
 val:
-  _target_: diffulab.datasets.ImageNetLatentREPA
+  _target_: diffulab.datasets.ImageNetLatent
   data_path: "data/imagenet"
   local: true
   batch_size: 128

configs/dataset/imagenet_noised_latents.yaml

@@ -0,0 +1,16 @@
+# CIFAR10 dataset configuration
+train:
+  _target_: diffulab.datasets.ImageNetNoisyLatent
+  data_path: "data/imagenet"
+  local: true
+  batch_size: 128
+  split: "train"
+  noise_tau: 0.8
+
+val:
+  _target_: diffulab.datasets.ImageNetNoisyLatent
+  data_path: "data/imagenet"
+  local: true
+  batch_size: 128
+  split: "val"
+  noise_tau : 0

configs/discriminator/rae.yaml

@@ -0,0 +1,3 @@
+_target_: diffulab.networks.disc.RAEDiscriminator
+model_name: facebook/dino-vits8
+features_depth: [2, 5, 8, 11]

configs/optimizer/adam.yaml

@@ -0,0 +1,5 @@
+_target_: torch.optim.Adam
+lr: 1e-4
+weight_decay: 0
+betas: [0.9, 0.999]
+eps: 1e-8

configs/train_imagenet_flow_matching_repa.yaml

@@ -1,10 +1,10 @@
-# configs/train_cifar10_flow_matching.yaml
+# configs/train_imagenet_flow_matching_repa.yaml
 # @package _global_
 defaults:
   - model: dit
   - diffuser: rectified_flow
   - trainer: default
-  - dataset: imagenet_repa
+  - dataset: imagenet_latents
   - dataloader: default
   - optimizer: adamw
   - vision_tower: dcae

configs/train_imagenet_rae_decoder.yaml

@@ -0,0 +1,41 @@
+defaults:
+  - vision_tower: rae
+  - discriminator: rae
+  - dataloader: default
+  - dataset: imagenet_noised_latents
+  - /optimizer@optimizer.rae_decoder: adam
+  - /optimizer@optimizer.discriminator: adam
+  - trainer: default
+
+vision_tower:
+  load_encoder: false
+
+optimizer:
+  rae_decoder:
+    betas: [0.5, 0.9]
+    lr : 2e-4
+  discriminator:
+    betas: [0.5, 0.9]
+    lr : 2e-4
+
+dataloader:
+  batch_size: 32
+
+dataset:
+  train:
+    batch_size: 32
+  val:
+    batch_size: 32
+
+trainer:
+  project_name: imagenet_rae_decoder
+  n_epoch: 16
+  precision_type: "bf16"
+  per_batch_scheduler: true
+  disc_epoch_start: 6
+  gan_epoch_start: 8
+  lpips_epoch_start: 1
+  lambda_lpips: 1
+  lambda_gan: 0.75
+  use_adaptive_weight_loss: true
+  gradient_accumulation_step: 16

configs/vision_tower/rae.yaml

@@ -0,0 +1,19 @@
+_target_: diffulab.networks.vision_towers.RAE
+decoder:
+  _target_: diffulab.networks.vision_towers.RAEDecoder
+  out_size: [256, 256]
+  out_channels: 3
+  encoder_dim: 768
+  input_dim: 1024
+  hidden_dim: 1024
+  num_heads: 16
+  mlp_ratio: 4
+  patch_size: 16
+  depth: 24
+  partial_rotary_factor: 1
+  use_checkpoint: False
+  dropout_attn: 0
+  dropout_mlp: 0
+dinov3_id: facebook/dinov3-vitb16-pretrain-lvd1689m
+load_encoder: true
+encoder_patch_size: 16

examples/train_diffusion.py

@@ -5,7 +5,7 @@
 from torch.utils.data import DataLoader
 
 from diffulab.diffuse import Diffuser
-from diffulab.training import Trainer
+from diffulab.training import BaseTrainer
 
 
 @hydra.main(version_base=None, config_path="../configs", config_name="train_mnist_flow_matching")
@@ -52,8 +52,7 @@ def count_parameters(model: torch.nn.Module) -> int:
         params=denoiser.parameters(),
     )
 
-    # TODO: add a run name for wandb
-    trainer = Trainer(
+    trainer = BaseTrainer(
         n_epoch=cfg.trainer.n_epoch,
         gradient_accumulation_step=cfg.trainer.gradient_accumulation_step,
         precision_type=cfg.trainer.precision_type,
@@ -62,6 +61,10 @@ def count_parameters(model: torch.nn.Module) -> int:
         ema_update_after_step=cfg.trainer.get("ema_update_after_step", 0),
         ema_update_every=cfg.trainer.get("ema_update_every", 10),
         run_config=OmegaConf.to_container(cfg, resolve=True),  # type: ignore[reportArgumentType]
+        compile=cfg.trainer.get("compile", False),
+        init_kwargs={
+            "wandb": cfg.trainer.get("wandb", {}),
+        },
     )
 
     trainer.train(

examples/train_rae_decoder.py

@@ -0,0 +1,125 @@
+import math
+
+import hydra
+import torch
+from hydra.utils import instantiate
+from omegaconf import DictConfig, OmegaConf
+from torch.optim.lr_scheduler import LambdaLR
+from torch.utils.data import DataLoader
+
+from diffulab.datasets.imagenet import ImageNetLatent
+from diffulab.networks.disc import RAEDiscriminator
+from diffulab.networks.vision_towers.rae import RAE
+from diffulab.training.trainers.extra.rae_trainer import RAETrainer
+
+
+def cosine_with_warmup_and_min_lr_lambda(
+    current_step: int, num_warmup_steps: int, num_training_steps: int, min_lr_factor: float = 0.1
+) -> float:
+    if current_step < num_warmup_steps:
+        return float(current_step) / float(max(1, num_warmup_steps))
+    progress = (current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
+    # interpolate between 1 and min_lr_factor
+    return min_lr_factor + (1 - min_lr_factor) * cosine
+
+
+def get_cosine_schedule_with_warmup_and_min_lr(
+    optimizer: torch.optim.Optimizer, num_warmup_steps: int, num_training_steps: int, min_lr_factor: float = 0.1
+) -> LambdaLR:
+    lr_lambda = lambda step: cosine_with_warmup_and_min_lr_lambda(  # type: ignore
+        step,  # type: ignore
+        num_warmup_steps,
+        num_training_steps,
+        min_lr_factor,  # type: ignore
+    )
+    return LambdaLR(optimizer, lr_lambda)  # type: ignore
+
+
+@hydra.main(version_base=None, config_path="../configs", config_name="train_imagenet_rae_decoder")
+def train(cfg: DictConfig) -> None:
+    print(OmegaConf.to_yaml(cfg))
+
+    train_dataset: ImageNetLatent = instantiate(cfg.dataset.train)
+    val_dataset: ImageNetLatent = instantiate(cfg.dataset.val)
+
+    train_dataset.set_latent_scale(1)
+    val_dataset.set_latent_scale(1)
+
+    dl_cfg = cfg.get("dataloader", {})
+    train_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=dl_cfg.get("batch_size", 32),
+        shuffle=dl_cfg.get("shuffle", True),
+        num_workers=dl_cfg.get("num_workers", 0),
+        pin_memory=dl_cfg.get("pin_memory", False),
+    )
+
+    val_loader = DataLoader(
+        dataset=val_dataset,
+        batch_size=dl_cfg.get("batch_size", 32),
+        shuffle=dl_cfg.get("shuffle", False),
+        num_workers=dl_cfg.get("num_workers", 0),
+        pin_memory=dl_cfg.get("pin_memory", False),
+    )
+
+    rae: RAE = instantiate(cfg.vision_tower)
+    discriminator: RAEDiscriminator = instantiate(cfg.discriminator)
+
+    def count_parameters(model: torch.nn.Module) -> int:
+        return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    print(f"Number of trainable parameters in the rae decoder: {count_parameters(rae.decoder):,}")
+    print(f"Number of trainable parameters in the discriminator: {count_parameters(discriminator):,}")
+
+    rae_optimizer = instantiate(cfg.optimizer.rae_decoder, params=rae.decoder.parameters())
+    disc_optimizer = instantiate(cfg.optimizer.discriminator, params=discriminator.parameters())
+
+    rae_scheduler = get_cosine_schedule_with_warmup_and_min_lr(
+        optimizer=rae_optimizer,
+        num_warmup_steps=len(train_loader),
+        num_training_steps=cfg.trainer.n_epoch * len(train_loader),
+    )
+    disc_scheduler = get_cosine_schedule_with_warmup_and_min_lr(
+        optimizer=disc_optimizer,
+        num_warmup_steps=len(train_loader),
+        num_training_steps=(cfg.trainer.n_epoch - cfg.trainer.disc_epoch_start) * len(train_loader),
+    )
+
+    rae_trainer = RAETrainer(
+        n_epoch=cfg.trainer.n_epoch,
+        gradient_accumulation_step=cfg.trainer.gradient_accumulation_step,
+        precision_type=cfg.trainer.precision_type,
+        project_name=cfg.trainer.project_name,
+        use_ema=cfg.trainer.use_ema,
+        ema_update_after_step=cfg.trainer.get("ema_update_after_step", 0),
+        ema_update_every=cfg.trainer.get("ema_update_every", 10),
+        run_config=OmegaConf.to_container(cfg, resolve=True),  # type: ignore[reportArgumentType]
+        compile=cfg.trainer.get("compile", False),
+        init_kwargs={
+            "wandb": cfg.trainer.get("wandb", {}),
+        },
+    )
+
+    rae_trainer.train(
+        rae=rae,
+        disc=discriminator,
+        rae_optimizer=rae_optimizer,
+        disc_optimizer=disc_optimizer,
+        train_dataloader=train_loader,
+        val_dataloader=val_loader,
+        rae_scheduler=rae_scheduler,
+        disc_scheduler=disc_scheduler,
+        per_batch_scheduler=cfg.trainer.get("per_batch_scheduler", True),
+        log_validation_images=cfg.trainer.get("log_validation_images", True),
+        disc_epoch_start=cfg.trainer.get("disc_epoch_start", 6),
+        gan_epoch_start=cfg.trainer.get("gan_epoch_start", 8),
+        lpips_epoch_start=cfg.trainer.get("lpips_epoch_start", 1),
+        lambda_lpips=cfg.trainer.get("lambda_lpips", 1.0),
+        lambda_gan=cfg.trainer.get("lambda_gan", 0.75),
+        use_adaptive_weight_loss=cfg.trainer.get("use_adaptive_weight_loss", True),
+    )
+
+
+if __name__ == "__main__":
+    train()

examples/train_repa.py

@@ -7,7 +7,7 @@
 from torch.utils.data import DataLoader
 
 from diffulab.diffuse import Diffuser
-from diffulab.training import Trainer
+from diffulab.training import BaseTrainer
 from diffulab.training.losses.repa import RepaLoss
 
 
@@ -77,7 +77,7 @@ def count_parameters(model: torch.nn.Module) -> int:
         + list(repa_loss.resampler.parameters() if repa_loss.resampler else []),
     )
 
-    trainer = Trainer(
+    trainer = BaseTrainer(
         n_epoch=cfg.trainer.n_epoch,
         gradient_accumulation_step=cfg.trainer.gradient_accumulation_step,
         precision_type=cfg.trainer.precision_type,