Add dataloader and add logging for training script

data/mnist_dataset.py

@@ -0,0 +1,71 @@
+import glob
+import os
+
+import cv2
+import numpy as np
+import torchvision
+from PIL import Image
+from torch.utils.data.dataset import Dataset
+from tqdm import tqdm
+
+
+class MnistDataset(Dataset):
+    r"""
+    Nothing special here. Just a simple dataset class for mnist images.
+    Created a dataset class rather using torchvision to allow
+    replacement with any other image dataset
+    """
+
+    def __init__(self, split, im_path, im_ext="png", im_size=28, return_hints=False):
+        r"""
+        Init method for initializing the dataset properties
+        :param split: train/test to locate the image files
+        :param im_path: root folder of images
+        :param im_ext: image extension. assumes all
+        images would be this type.
+        """
+        self.split = split
+        self.im_ext = im_ext
+        self.return_hints = return_hints
+        self.images = self.load_images(im_path)
+
+    def load_images(self, im_path):
+        r"""
+        Gets all images from the path specified
+        and stacks them all up
+        :param im_path:
+        :return:
+        """
+        assert os.path.exists(im_path), "images path {} does not exist".format(im_path)
+        ims = []
+        labels = []
+        for d_name in tqdm(os.listdir(im_path)):
+            for fname in glob.glob(
+                os.path.join(im_path, d_name, "*.{}".format(self.im_ext))
+            ):
+                ims.append(fname)
+        print("Found {} images for split {}".format(len(ims), self.split))
+        return ims
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, index):
+        im = Image.open(self.images[index])
+        im_tensor = torchvision.transforms.ToTensor()(im)
+
+        # Convert input to -1 to 1 range.
+        im_tensor = (2 * im_tensor) - 1
+
+        if self.return_hints:
+            canny_image = Image.open(self.images[index])
+            canny_image = np.array(canny_image)
+            canny_image = cv2.Canny(canny_image, 100, 200)
+            canny_image = canny_image[:, :, None]
+            canny_image = np.concatenate(
+                [canny_image, canny_image, canny_image], axis=2
+            )
+            canny_image_tensor = torchvision.transforms.ToTensor()(canny_image)
+            return im_tensor, canny_image_tensor
+        else:
+            return im_tensor

model_config/mnist.yaml

@@ -0,0 +1,35 @@
+dataset_params:
+  im_path: 'data/mnist/train/images'
+  im_test_path: 'data/mnist/test/images'
+  canny_im_size: 28
+
+diffusion_params:
+  num_timesteps : 1000
+  beta_start : 0.0001
+  beta_end : 0.02
+
+model_params:
+  im_channels : 1
+  im_size : 28
+  hint_channels : 3
+  down_channels : [32, 64, 128, 256]
+  mid_channels : [256, 256, 128]
+  down_sample : [True, True, False]
+  time_emb_dim : 128
+  num_down_layers : 2
+  num_mid_layers : 2
+  num_up_layers : 2
+  num_heads : 4
+
+train_params:
+  task_name: 'mnist'
+  batch_size: 64
+  num_epochs: 40
+  controlnet_epochs : 1
+  num_samples : 25
+  num_grid_rows : 5
+  save_epoch: 2
+  ddpm_lr: 0.0001
+  controlnet_lr: 0.0001
+  ddpm_ckpt_name: 'ddpm_ckpt.pth'
+  controlnet_ckpt_name: 'ddpm_controlnet_ckpt.pth'

training_scripts/train_ddpm.py

@@ -0,0 +1,305 @@
+import logging
+import os
+
+import numpy as np
+import torch
+import wandb
+import yaml
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from data import mnist_dataset
+from data.mnist_dataset import MnistDataset
+from model_blocks.unet_base import UNet
+from scheduler.linear_scheduler import LinearNoiseScheduler
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+logger = logging.getLogger(__name__)
+wandb.login()
+
+
+def init_wandb(config):
+    """
+    Initialize a new wandb run
+    """
+    run = wandb.init(
+        project="controlnet-ddpm-mnist",
+        config=config,
+        resume="allow",  # Allows resuming if run was interrupted
+    )
+    return run
+
+
+def load_checkpoint(model, optimizer, scheduler, checkpoint_path):
+    """
+    Load model checkpoint from local file
+    """
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+
+    model.load_state_dict(checkpoint["model_state_dict"])
+    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+    if (
+        scheduler
+        and "scheduler_state_dict" in checkpoint
+        and checkpoint["scheduler_state_dict"]
+    ):
+        scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
+
+    start_epoch = checkpoint["epoch"] + 1  # Start from the next epoch
+    step = checkpoint["step"]
+
+    print(f"Loaded checkpoint from epoch {checkpoint['epoch']}")
+    return start_epoch, step
+
+
+def save_checkpoint(
+    model, optimizer, scheduler, epoch, loss, step, run, checkpoint_path
+):
+    """
+    Save model checkpoint locally and to wandb
+    """
+    # Create checkpoint dictionary
+    checkpoint = {
+        "epoch": epoch,
+        "model_state_dict": model.state_dict(),
+        "optimizer_state_dict": optimizer.state_dict(),
+        "scheduler_state_dict": scheduler.state_dict() if scheduler else None,
+        "loss": loss,
+        "step": step,
+    }
+
+    # Save locally
+    torch.save(checkpoint, checkpoint_path)
+
+    # Log to wandb
+    artifact = wandb.Artifact(f"model-checkpoint-epoch-{epoch}", type="model")
+    artifact.add_file(checkpoint_path)
+    run.log_artifact(artifact)
+
+    print(f"Checkpoint saved at epoch {epoch}")
+    return checkpoint_path
+
+
+def train(args):
+    with open(args.config_path, "r") as file:
+        try:
+            config = yaml.safe_load(file)
+        except yaml.YAMLError as exc:
+            print(exc)
+    print(config)
+    run = init_wandb(config["train_params"])
+
+    diffusion_config = config["diffusion_params"]
+    dataset_config = config["dataset_params"]
+    model_config = config["model_params"]
+    train_config = config["train_params"]
+    scheduler = LinearNoiseScheduler(
+        num_timesteps=diffusion_config["num_timesteps"],
+        beta_start=diffusion_config["beta_start"],
+        beta_end=diffusion_config["beta_end"],
+    )
+
+    mnist = MnistDataset("train", dataset_config["im_path"])
+    mnist_loader = DataLoader(
+        mnist, batch_size=train_config["batch_size"], shuffle=True, num_workers=4
+    )
+
+    model = UNet(model_config).to(device)
+    model.train()
+    logger.debug(f"Initialized model and set to train")
+    optimizer = torch.optim.AdamW(model.parameters(), lr=train_config["ddpm_lr"])
+    criterion = torch.nn.MSELoss()
+
+    # Create the output directories
+    if not os.path.exists(
+        os.path.join(train_config["task_name"], train_config["ddpm_ckpt_name"])
+    ):
+        os.mkdir(train_config["task_name"])
+
+    ckpt_path = os.path.join(train_config["task_name"], train_config["ddpm_ckpt_name"])
+    # Load checkpoint if there
+    if os.path.exists(ckpt_path):
+        start_epoch, step = load_checkpoint(
+            model, optimizer=optimizer, scheduler=scheduler, checkpoint_path=ckpt_path
+        )
+    else:
+        start_epoch = 0
+        step = 0
+
+    # Log model architecture as a Table
+    model_table = wandb.Table(columns=["Layer", "Parameters"])
+    total_params = 0
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            params = param.numel()
+            total_params += params
+            model_table.add_data(name, params)
+
+    wandb.log({"model_architecture": model_table})
+    wandb.log({"total_parameters": total_params})
+
+    # Watch model gradients and parameters
+    wandb.watch(model, log="all", log_freq=100)
+
+    for epoch in range(start_epoch, train_config["num_epochs"]):
+        losses = []
+        batch_idx = 0
+        progress_bar = tqdm(
+            mnist_loader, desc=f"Epoch {epoch + 1}/{train_config['num_epochs']}"
+        )
+
+        for im in progress_bar:
+            batch_idx += 1
+            optimizer.zero_grad()
+            im = im.float().to(device)
+
+            # Sample noise
+            noise = torch.randn_like(im).to(device)
+            logger.debug(f"Sampled noise epoch {epoch} : {noise.shape}")
+
+            # Sample timestep
+            t = torch.randint(0, diffusion_config["num_timesteps"], (im.shape[0],)).to(
+                device
+            )
+
+            noisy_im = scheduler.add_noise(im, noise, t)
+            noise_pred = model(noisy_im, t)
+
+            loss = criterion(noise_pred, noise)
+            losses.append(loss.item())
+            loss.backward()
+            optimizer.step()
+
+            # Calculate gradient norm for monitoring
+            total_norm = 0
+            for p in model.parameters():
+                if p.grad is not None:
+                    param_norm = p.grad.data.norm(2)
+                    total_norm += param_norm.item() ** 2
+            total_norm = total_norm**0.5
+
+            # Update progress bar
+            progress_bar.set_postfix({"loss": loss.item(), "avg_loss": np.mean(losses)})
+
+            wandb.log(
+                {
+                    "train/batch_loss": loss.item(),
+                    "train/step": step,
+                    "train/epoch": epoch + batch_idx / len(mnist_loader),
+                    "train/gradient_norm": total_norm,
+                    "train/learning_rate": optimizer.param_groups[0]["lr"],
+                }
+            )
+
+            step += 1
+
+        avg_loss = np.mean(losses)
+
+        # Log epoch-level metrics
+        wandb.log(
+            {
+                "train/epoch_loss": avg_loss,
+                "train/epoch_completed": epoch,
+            }
+        )
+
+        print(f"Finished epoch: {epoch} | Loss: {np.mean(losses):.4f}")
+        if epoch % train_config["save_epoch"]:
+            visualize_samples(
+                model, scheduler, epoch, diffusion_config["num_timesteps"], device
+            )
+
+        save_checkpoint(
+            model, optimizer, scheduler, epoch, np.mean(losses), step, run, ckpt_path
+        )
+
+    # Log final model as artifact
+    logging.info("Finished training and starting to save model")
+    final_model_path = os.path.join(
+        train_config["task_name"], f"final_{train_config['ddpm_ckpt_name']}"
+    )
+    save_checkpoint(
+        model,
+        optimizer,
+        scheduler,
+        train_config["num_epochs"] - 1,
+        avg_loss,
+        step,
+        run,
+        final_model_path,
+    )
+    logging.info("Saved Model to Wandb and local")
+
+    # Log a summary table of training
+    summary_table = wandb.Table(columns=["Metric", "Value"])
+    summary_table.add_data("Final Loss", avg_loss)
+    summary_table.add_data("Best Loss", best_loss)
+    summary_table.add_data("Best Epoch", epoch)
+    summary_table.add_data("Total Steps", step)
+    summary_table.add_data("Training Time (hours)", wandb.run.duration / 3600)
+
+    wandb.log({"training_summary": summary_table})
+
+    # Finish the run
+    wandb.finish()
+
+
+def visualize_samples(model, scheduler, epoch, num_timesteps, device, num_samples=4):
+    """
+    Generate sample images from noise and log to wandb
+    """
+    model.eval()
+    with torch.no_grad():
+        # Start with random noise
+        samples = torch.randn(num_samples, 1, 28, 28).to(device)
+
+        # Store the denoising process
+        sample_images = []
+
+        # Record more frequently at the beginning of sampling
+        log_steps = set([0, 20, 50, 100, 200, 400, 600, 800, num_timesteps - 1])
+
+        # Denoise gradually
+        for i in tqdm(reversed(range(num_timesteps)), desc="Sampling"):
+            t = torch.full((num_samples,), i, device=device, dtype=torch.long)
+
+            # Get model prediction and update sample
+            predicted_noise = model(samples, t)
+            samples = scheduler.step(predicted_noise, i, samples)
+
+            # Save images at specified timesteps
+            if i in log_steps:
+                # Denormalize and convert to numpy for logging
+                denorm_samples = samples.clamp(-1, 1).cpu().numpy()
+                denorm_samples = (
+                    denorm_samples + 1
+                ) / 2.0  # scale from [-1, 1] to [0, 1]
+                sample_images.append((i, denorm_samples))
+
+    # Create a grid to show denoising process
+    images_to_log = {}
+
+    # Log individual samples
+    for i, sample in enumerate(samples):
+        sample_np = sample.clamp(-1, 1).cpu().numpy()
+        sample_np = (sample_np + 1) / 2.0  # scale from [-1, 1] to [0, 1]
+        images_to_log[f"sample_{i}_epoch_{epoch}"] = wandb.Image(
+            sample_np[0], caption=f"Sample {i}, Epoch {epoch}"
+        )
+
+    # Log denoising process for first sample
+    denoising_steps = []
+    for step_idx, samples_np in sample_images:
+        denoising_steps.append(
+            wandb.Image(
+                samples_np[0][0],
+                caption=f"Step {num_timesteps - step_idx}/{num_timesteps}",
+            )
+        )
+
+    images_to_log["denoising_process_epoch_" + str(epoch)] = denoising_steps
+
+    # Log all images
+    wandb.log(images_to_log)
+
+    model.train()