Add the training script for dit

celeba/config.yaml

@@ -11,10 +11,10 @@ diffusion_params:
 dit_params:
   patch_size: 2
   num_layers: 12
-  hidden_size: 768
+  hidden_dim: 768
   num_heads: 12
   head_dim: 64
-  timestep_emb_dim: 768
+  temb_dim: 768
 
 autoencoder_params:
   z_channels: 4

train_dit.py

@@ -0,0 +1,128 @@
+import argparse
+import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+import yaml
+from torch.optim import AdamW
+from tqdm import tqdm
+
+from celeba import create_dataloader
+from model.transformer import DIT
+from model.vae import VAE
+from scheduler.linear_scheduler import LinearNoiseScheduler
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def train(args):
+    with open(args.config_path, "r") as file:
+        try:
+            config = yaml.safe_load(file)
+        except yaml.YAMLError as e:
+            print(f"Error in loading yaml: {e}")
+
+    train_config = config["train_params"]
+    dit_config = config["dit_params"]
+    dataset_config = config["dataset_params"]
+    diffusion_params = config["diffusion_params"]
+    vae_config = config["autoencoder_params"]
+
+    dataloader = create_dataloader(dataset_config["im_path"])
+
+    scheduler = LinearNoiseScheduler(
+        diffusion_params["num_timesteps"],
+        diffusion_params["beta_start"],
+        diffusion_params["beta_end"],
+    )
+
+    im_size = dataset_config["im_size"] // 2 ** sum(vae_config["down_sample"])
+    model = DIT(
+        im_size=im_size, im_channels=dataset_config["im_channels"], config=dit_config
+    ).to(device)
+    model.train()
+
+    if os.path.exists(
+        os.path.join(train_config["task_name"], train_config["dit_ckpt_name"])
+    ):
+        checkpoint = torch.load(
+            os.path.join(train_config["task_name"], train_config["dit_ckpt_name"]),
+            map_location=device,
+        )
+
+        model.load_state_dict(checkpoint["dit"])
+        start_epoch = checkpoint["epoch"]
+        step_count = checkpoint["step_count"]
+    else:
+        step_count = 0
+        start_epoch = 0
+
+    if not os.path.exists(
+        os.path.join(
+            train_config["task_name"], train_config["vae_autoencoder_ckpt_name"]
+        )
+    ):
+        print("No VAE  checkpoint found, VAE checkpoint needed")
+        return
+    else:
+        vae = VAE(dataset_config["im_channels"], vae_config).to(device)
+        vae.load_state_dict(
+            torch.load(
+                os.path.join(
+                    train_config["task_name"], train_config["vae_autoencoder_ckpt_name"]
+                ),
+                map_location=device,
+            )
+        )
+        vae.eval()
+        for param in vae.parameters():
+            param.requires_grad = False
+        print("VAE checkpoint loaded")
+
+    num_epochs = train_config["dit_epochs"]
+    optimizer = AdamW(model.parameters(), lr=train_config["dit_lr"])
+    accu_steps = train_config["dit_acc_steps"]
+    criterion = nn.MSELoss()
+
+    for epoch in range(start_epoch, num_epochs):
+        losses = []
+        for im in tqdm(dataloader):
+            im = im.float().to(device)
+            step_count += 1
+            with torch.no_grad():
+                im, _ = vae.encode(im)
+
+            noise = torch.randn_like(im).to(device)
+
+            t = torch.randint(0, diffusion_params["num_time_steps"], (im.shape[0],)).to(
+                device
+            )
+            noisy_im = scheduler.add_noise(im, noise, t)
+            pred = model(noisy_im, t)
+            loss = criterion(pred, noise)
+            losses.append(loss.item())
+            loss = loss / accu_steps
+            loss.backward()
+            if step_count % accu_steps == 0:
+                optimizer.step()
+                optimizer.zero_grad()
+        optimizer.step()
+        optimizer.zero_grad()
+        print(f"Epoch {epoch}: Loss: {np.mean(losses)}")
+        torch.save(
+            {"dit": model.state_dict(), "epoch": epoch + 1, "step": step_count},
+            os.path.join(
+                train_config["task_name"], train_config["vae_autoencoder_ckpt_name"]
+            ),
+        )
+    print("Done Training")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Arguments for dit training")
+    parser.add_argument(
+        "--config", dest="config_path", default="celeba/config.yaml", type=str
+    )
+    args = parser.parse_args()
+    train(args)