diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/__init__.py | 50 | ||||
| -rw-r--r-- | src/config.py | 34 | ||||
| -rw-r--r-- | src/dataloader.py | 210 | ||||
| -rw-r--r-- | src/diffusion.py | 154 | ||||
| -rw-r--r-- | src/metrics.py | 98 | ||||
| -rw-r--r-- | src/sample.py | 248 | ||||
| -rw-r--r-- | src/train_autoencoder.py | 167 | ||||
| -rw-r--r-- | src/train_diffusion.py | 197 | ||||
| -rw-r--r-- | src/utils.py | 110 |
9 files changed, 1268 insertions, 0 deletions
diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..220bb0e --- /dev/null +++ b/src/__init__.py @@ -0,0 +1,50 @@ +""" +MNIST Diffusion Model package. +""" + +from .config import ( + DEVICE, + IMG_SIZE, + NB_CHANNEL, + NB_LABEL, + DIFFU_STEPS, + ALPHA, + ALPHA_BAR, + BETA, + EPOCHS, + BATCH_SIZE, + LEARNING_RATE, + CHECKPOINT_DIR, + PLOTS_DIR, +) + +from .diffusion import ( + q_xt_xt_1, + q_xt_x0, + p_xt_1_xt, + p_xt_1_xt_x0_pred, + forward_diffusion, +) + +from .dataloader import ( + get_mnist_dataset, + get_diffusion_dataloader, + get_autoencoder_dataloader, + DiffusionDataset, + DiffusionDatasetX0, + AutoencoderDataset, +) + +from .utils import ( + tensor_to_image, + tensor_to_images, + save_checkpoint, + load_checkpoint, + save_plot, +) + +from .metrics import ( + fid, + kl_divergence, + jsd, +) diff --git a/src/config.py b/src/config.py new file mode 100644 index 0000000..4ac56c4 --- /dev/null +++ b/src/config.py @@ -0,0 +1,34 @@ +""" +Configuration file for MNIST Diffusion model. +Contains all hyperparameters and constants. +""" + +import torch + +# Device configuration +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +# Image parameters (MNIST) +IMG_SIZE = 28 +NB_CHANNEL = 1 +NB_LABEL = 10 + +# Diffusion parameters +DIFFU_STEPS = 1000 + +# Noise schedule (linear beta schedule) +BETA = torch.linspace(1e-4, 2e-2, DIFFU_STEPS, device=DEVICE) +BETA = torch.cat((torch.tensor([0.0], device=DEVICE), BETA)) +ALPHA = 1 - BETA +ALPHA_BAR = torch.cumprod(ALPHA, dim=0) + +# Training parameters +EPOCHS = 10 +BATCH_SIZE = 64 +LEARNING_RATE = 2e-4 +NUM_WORKERS = 4 + +# Paths +DATA_DIR = "data" +CHECKPOINT_DIR = "checkpoints" +PLOTS_DIR = "plots" diff --git a/src/dataloader.py b/src/dataloader.py new file mode 100644 index 0000000..a08ae27 --- /dev/null +++ b/src/dataloader.py @@ -0,0 +1,210 @@ +""" +Data loading utilities for MNIST diffusion training. +""" + +import torch +from torch.utils.data import Dataset, DataLoader +from torchvision import datasets, transforms + +from .config import DEVICE, DIFFU_STEPS, NB_LABEL, DATA_DIR +from .diffusion import q_xt_x0 + + +def get_mnist_dataset(train: bool = True) -> datasets.MNIST: + """ + Load MNIST dataset. + + Args: + train: If True, load training set. Otherwise load test set. + + Returns: + MNIST dataset + """ + return datasets.MNIST( + root=DATA_DIR, + train=train, + download=True, + transform=transforms.ToTensor(), + ) + + +class DiffusionDataset(Dataset): + """ + Dataset wrapper for diffusion training. + Returns noisy image, timestep, label, and target noise. + + Args: + dataset: Base image dataset (e.g., MNIST) + autoencoder: Optional autoencoder for latent diffusion + """ + + def __init__(self, dataset: Dataset, autoencoder: torch.nn.Module = None): + super().__init__() + self.dataset = dataset + self.autoencoder = autoencoder + + def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]: + # Get image and label + img, label = self.dataset[index] + img = img.to(DEVICE) + + # Optionally encode to latent space + if self.autoencoder is not None: + with torch.no_grad(): + img = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0) + + # Normalize to [-1, 1] + img = img * 2 - 1 + + # Sample random timestep and add noise + t = torch.randint(1, DIFFU_STEPS + 1, (1,), device=DEVICE) + xt, eps = q_xt_x0(img, t) + + # Convert label to one-hot vector + vec = torch.nn.functional.one_hot( + torch.tensor(min(label, NB_LABEL - 1)), + num_classes=NB_LABEL, + ) + + return ( + xt.clone().detach().to(dtype=torch.float32, device=DEVICE), + t.clone().detach().to(dtype=torch.float32, device=DEVICE), + vec.clone().detach().to(dtype=torch.float32, device=DEVICE), + eps, # Target: the noise that was added + ) + + def __len__(self) -> int: + return len(self.dataset) + + +class DiffusionDatasetX0(Dataset): + """ + Dataset wrapper for diffusion training where model predicts x0 instead of noise. + Returns noisy image, timestep, label, and target clean image. + + Args: + dataset: Base image dataset (e.g., MNIST) + autoencoder: Optional autoencoder for latent diffusion + """ + + def __init__(self, dataset: Dataset, autoencoder: torch.nn.Module = None): + super().__init__() + self.dataset = dataset + self.autoencoder = autoencoder + + def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]: + # Get image and label + img, label = self.dataset[index] + img = img.to(DEVICE) + + # Optionally encode to latent space + if self.autoencoder is not None: + with torch.no_grad(): + img = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0) + + # Normalize to [-1, 1] + img = img * 2 - 1 + + # Sample random timestep and add noise + t = torch.randint(1, DIFFU_STEPS + 1, (1,), device=DEVICE) + xt, _ = q_xt_x0(img, t) + + # Convert label to one-hot vector + vec = torch.nn.functional.one_hot( + torch.tensor(min(label, NB_LABEL - 1)), + num_classes=NB_LABEL, + ) + + return ( + xt.clone().detach().to(dtype=torch.float32, device=DEVICE), + t.clone().detach().to(dtype=torch.float32, device=DEVICE), + vec.clone().detach().to(dtype=torch.float32, device=DEVICE), + img.clone().detach().to(dtype=torch.float32, device=DEVICE), # Target: clean image + ) + + def __len__(self) -> int: + return len(self.dataset) + + +class AutoencoderDataset(Dataset): + """ + Dataset wrapper for autoencoder training. + Returns image as both input and target. + + Args: + dataset: Base image dataset + """ + + def __init__(self, dataset: Dataset): + self.dataset = dataset + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: + data = self.dataset[idx][0].to(DEVICE) + return data, data + + +def get_diffusion_dataloader( + predict_x0: bool = True, + batch_size: int = 64, + shuffle: bool = True, + num_workers: int = 4, + autoencoder: torch.nn.Module = None, +) -> DataLoader: + """ + Create a DataLoader for diffusion training. + + Args: + predict_x0: If True, model predicts x0. Otherwise predicts noise. + batch_size: Batch size + shuffle: Whether to shuffle data + num_workers: Number of data loading workers + autoencoder: Optional autoencoder for latent diffusion + + Returns: + DataLoader for training + """ + mnist = get_mnist_dataset(train=True) + + if predict_x0: + dataset = DiffusionDatasetX0(mnist, autoencoder) + else: + dataset = DiffusionDataset(mnist, autoencoder) + + return DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + persistent_workers=True if num_workers > 0 else False, + ) + + +def get_autoencoder_dataloader( + batch_size: int = 64, + shuffle: bool = True, + num_workers: int = 4, +) -> DataLoader: + """ + Create a DataLoader for autoencoder training. + + Args: + batch_size: Batch size + shuffle: Whether to shuffle data + num_workers: Number of data loading workers + + Returns: + DataLoader for training + """ + mnist = get_mnist_dataset(train=True) + dataset = AutoencoderDataset(mnist) + + return DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + num_workers=num_workers, + persistent_workers=True if num_workers > 0 else False, + ) diff --git a/src/diffusion.py b/src/diffusion.py new file mode 100644 index 0000000..0d63ee1 --- /dev/null +++ b/src/diffusion.py @@ -0,0 +1,154 @@ +""" +Diffusion process utilities. +Contains forward and reverse diffusion functions. +""" + +import torch +from .config import DEVICE, ALPHA, ALPHA_BAR, BETA, DIFFU_STEPS + + +def q_xt_xt_1(xt_1: torch.Tensor, t: int | torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Forward diffusion step: q(x_t | x_{t-1}). + Adds noise to image at step t-1 to get image at step t. + + Args: + xt_1: Image at timestep t-1 [B, C, H, W] + t: Timestep (int or tensor) + + Returns: + xt: Noisy image at timestep t + eps: The noise that was added + """ + if isinstance(t, int): + t_ind = torch.tensor(t, dtype=torch.long, device=DEVICE) + else: + t_ind = t.to(dtype=torch.long, device=DEVICE) + + alpha = ALPHA[t_ind] + mean = torch.sqrt(alpha) * xt_1 + std = torch.sqrt(1 - alpha) + + eps = torch.randn(xt_1.shape, device=DEVICE) + xt = mean + std * eps + + return xt, eps + + +def q_xt_x0(x0: torch.Tensor, t: int | torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Forward diffusion: q(x_t | x_0). + Directly compute noisy image at any timestep t from clean image x0. + + Args: + x0: Clean image [B, C, H, W] or [C, H, W] + t: Timestep tensor [B, 1] or [B] + + Returns: + xt: Noisy image at timestep t + eps: The noise that was added + """ + if isinstance(t, int): + t_ind = torch.tensor(t, dtype=torch.long, device=DEVICE) + else: + t_ind = t.to(dtype=torch.long, device=DEVICE) + + # Handle both batched and single images + reshaped = len(x0.shape) == 3 + if reshaped: + c, w, h = x0.shape + b = 1 + x0 = x0.view(b, c, w, h) + else: + b, c, w, h = x0.shape + + # Reshape t for broadcasting + t_ind = t_ind.view(b, 1, 1, 1) + t_ind = t_ind.expand(b, c, w, h) + + alpha_bar = ALPHA_BAR[t_ind] + mean = torch.sqrt(alpha_bar) * x0 + std = torch.sqrt(1 - alpha_bar) + + eps = torch.randn(x0.shape, device=DEVICE) + xt = mean + std * eps + + if reshaped: + xt = xt.view(c, w, h) + + return xt, eps + + +def p_xt_1_xt(model: torch.nn.Module, xt: torch.Tensor, t: torch.Tensor, + vec: torch.Tensor) -> torch.Tensor: + """ + Reverse diffusion step: p(x_{t-1} | x_t). + Denoise image at step t to get image at step t-1. + Model predicts the noise. + + Args: + model: UNet model that predicts noise + xt: Noisy image at timestep t [B, C, H, W] + t: Timestep tensor [B, 1] + vec: Label one-hot vector [B, NB_LABEL] + + Returns: + xt_1: Denoised image at timestep t-1 + """ + t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t + + alpha_bar_t = ALPHA_BAR[t_ind].view(-1, 1, 1, 1) + alpha_bar_t_1 = ALPHA_BAR[t_ind - 1].view(-1, 1, 1, 1) + alpha_t = ALPHA[t_ind].view(-1, 1, 1, 1) + beta_t = BETA[t_ind].view(-1, 1, 1, 1) + + beta_tilde = (1 - alpha_bar_t_1) / (1 - alpha_bar_t) * beta_t + + # Model predicts the noise + epsilon_theta = model(xt, t, vec) + + sigma_theta = torch.sqrt(beta_tilde) + mu_theta = (xt - (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) * epsilon_theta) / torch.sqrt(alpha_t) + + # Don't add noise at t=1 + mask_t0 = (t > 1).to(dtype=torch.float32).view(-1, 1, 1, 1) + noise = torch.randn(xt.shape, device=DEVICE) * mask_t0 + + return mu_theta + sigma_theta * noise + + +def p_xt_1_xt_x0_pred(model: torch.nn.Module, xt: torch.Tensor, t: torch.Tensor, + vec: torch.Tensor) -> torch.Tensor: + """ + Reverse diffusion step where model predicts x0 directly. + + Args: + model: UNet model that predicts clean image x0 + xt: Noisy image at timestep t [B, C, H, W] + t: Timestep tensor [B, 1] + vec: Label one-hot vector [B, NB_LABEL] + + Returns: + xt_1: Denoised image at timestep t-1 + """ + x0_pred = model(xt, t, vec) + xt_1, _ = q_xt_x0(x0_pred, t - 1) + return xt_1 + + +def forward_diffusion(x0: torch.Tensor) -> list[torch.Tensor]: + """ + Run full forward diffusion process. + + Args: + x0: Clean image [C, H, W] or [B, C, H, W] + + Returns: + List of images at each timestep [x0, x1, ..., xT] + """ + x = x0.clone() + xs = [x] + for t in range(1, DIFFU_STEPS + 1): + x, _ = q_xt_xt_1(x, t) + xs.append(x) + return xs diff --git a/src/metrics.py b/src/metrics.py new file mode 100644 index 0000000..08eb3cd --- /dev/null +++ b/src/metrics.py @@ -0,0 +1,98 @@ +import numpy as np +import scipy + + +def fid(reals: np.ndarray, fakes: np.ndarray) -> float: + """ + Calculate Frechet Inception Distance (FID) score. + + Args: + reals: Real images [N, C, H, W] + fakes: Generated images [N, C, H, W] + + Returns: + FID score (lower is better) + """ + reals = reals.reshape(reals.shape[0], -1) + fakes = fakes.reshape(fakes.shape[0], -1) + + mu_real = np.mean(reals, axis=0) + mu_fake = np.mean(fakes, axis=0) + sigma_real = np.cov(reals, rowvar=False) + sigma_fake = np.cov(fakes, rowvar=False) + + diff = mu_real - mu_fake + covmean, _ = scipy.linalg.sqrtm(sigma_real.dot(sigma_fake), disp=False) + + if not np.isfinite(covmean).all(): + eps = 1e-6 + offset = np.eye(sigma_real.shape[0]) * eps + covmean = scipy.linalg.sqrtm((sigma_real + offset).dot(sigma_fake + offset)) + + if np.iscomplexobj(covmean): + covmean = covmean.real + + return diff @ diff + np.trace(sigma_real) + np.trace(sigma_fake) - 2 * np.trace(covmean) + + +def kl_divergence(reals: np.ndarray, fakes: np.ndarray) -> float: + """ + Calculate KL divergence between real and fake image distributions. + + Args: + reals: Real images [N, C, H, W] + fakes: Generated images [N, C, H, W] + + Returns: + KL divergence value + """ + reals = reals.transpose(1, 0, 2, 3).reshape(reals.shape[1], -1) + fakes = fakes.transpose(1, 0, 2, 3).reshape(fakes.shape[1], -1) + + hist_real = np.apply_along_axis( + lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, reals + ) + hist_fake = np.apply_along_axis( + lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, fakes + ) + + # Add smoothing + hist_real = hist_real + 1 + hist_fake = hist_fake + 1 + + hist_real = hist_real / np.sum(hist_real) + hist_fake = hist_fake / np.sum(hist_fake) + + return np.mean(np.log(hist_real / hist_fake)) + + +def jsd(reals: np.ndarray, fakes: np.ndarray) -> float: + """ + Calculate Jensen-Shannon divergence between real and fake image distributions. + + Args: + reals: Real images [N, C, H, W] + fakes: Generated images [N, C, H, W] + + Returns: + JSD value + """ + reals = reals.transpose(1, 0, 2, 3).reshape(reals.shape[1], -1) + fakes = fakes.transpose(1, 0, 2, 3).reshape(fakes.shape[1], -1) + + hist_real = np.apply_along_axis( + lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, reals + ) + hist_fake = np.apply_along_axis( + lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, fakes + ) + + hist_real = hist_real + 1 + hist_fake = hist_fake + 1 + + hist_real = hist_real / np.sum(hist_real) + hist_fake = hist_fake / np.sum(hist_fake) + + hist_avg = (hist_real + hist_fake) / 2 + + return 0.5 * (np.mean(np.log(hist_real / hist_avg)) + np.mean(np.log(hist_fake / hist_avg))) diff --git a/src/sample.py b/src/sample.py new file mode 100644 index 0000000..5b438ed --- /dev/null +++ b/src/sample.py @@ -0,0 +1,248 @@ +""" +Sampling and evaluation script for trained diffusion model. +Generate samples from a trained model and visualize results. +""" + +from models import UNetMNIST +from src.utils import ensure_dirs, tensor_to_image, load_checkpoint +from src.dataloader import get_mnist_dataset +from src.diffusion import p_xt_1_xt_x0_pred, forward_diffusion, q_xt_x0 +from src.config import ( + DEVICE, DIFFU_STEPS, NB_CHANNEL, IMG_SIZE, NB_LABEL, + CHECKPOINT_DIR, PLOTS_DIR +) +import os +import torch +import numpy as np +import matplotlib.pyplot as plt +from rich.progress import track + +import sys +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +def generate_samples( + model: torch.nn.Module, + n_samples: int = 10, + labels: list[int] | None = None, +) -> torch.Tensor: + """ + Generate samples from the diffusion model. + + Args: + model: Trained UNet model + n_samples: Number of samples per class (or total if labels provided) + labels: Optional list of specific labels to generate + + Returns: + Generated samples tensor [N, C, H, W] + """ + model.eval() + + if labels is None: + # Generate samples for all classes + labels = list(range(NB_LABEL)) * n_samples + + n_total = len(labels) + + with torch.no_grad(): + # Start from pure noise + x = torch.randn(n_total, NB_CHANNEL, IMG_SIZE, IMG_SIZE, device=DEVICE) + + # Create one-hot label vectors + vec = torch.tensor(labels, device=DEVICE) + vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(dtype=torch.float32) + + # Reverse diffusion process + for t in track(range(DIFFU_STEPS, 0, -1), description="Generating samples"): + t_tensor = torch.tensor([[t]] * n_total, device=DEVICE, dtype=torch.float32) + x = p_xt_1_xt_x0_pred(model, x, t_tensor, vec) + + # Normalize to [0, 1] + x = x * 0.5 + 0.5 + x = x.clamp(0, 1) + + return x + + +def visualize_forward_diffusion(save_path: str | None = None): + """ + Visualize the forward diffusion process on a real image. + + Args: + save_path: Path to save the visualization + """ + ensure_dirs() + + # Get a random image from MNIST + dataset = get_mnist_dataset(train=True) + idx = np.random.randint(0, len(dataset)) + img, label = dataset[idx] + img = img.to(DEVICE) + + # Normalize to [-1, 1] + img = img * 2 - 1 + + # Run forward diffusion + xs = forward_diffusion(img) + + # Select timesteps to visualize + n_plots = 10 + timesteps = np.linspace(1, DIFFU_STEPS, n_plots, dtype=int) + + fig, axes = plt.subplots(2, n_plots + 1, figsize=(2 * n_plots, 5)) + + # Row labels + axes[0, 0].text(0.5, 0.5, 'Step-by-step', ha='center', va='center', fontsize=10) + axes[0, 0].axis('off') + axes[1, 0].text(0.5, 0.5, 'Direct', ha='center', va='center', fontsize=10) + axes[1, 0].axis('off') + + # Plot step-by-step diffusion + for i, t in enumerate(timesteps): + axes[0, i + 1].imshow(tensor_to_image(xs[t]), cmap='gray') + axes[0, i + 1].set_title(f't={t}') + axes[0, i + 1].axis('off') + + # Plot direct diffusion for comparison + xt, _ = q_xt_x0(img, t) + axes[1, i + 1].imshow(tensor_to_image(xt), cmap='gray') + axes[1, i + 1].axis('off') + + fig.suptitle(f'Forward Diffusion Process (Label: {label})') + plt.tight_layout() + + if save_path is None: + save_path = os.path.join(PLOTS_DIR, 'forward_diffusion.png') + fig.savefig(save_path) + print(f"Saved forward diffusion visualization to {save_path}") + plt.close(fig) + + +def visualize_backward_diffusion(model: torch.nn.Module, save_path: str | None = None): + """ + Visualize the backward (reverse) diffusion process. + + Args: + model: Trained UNet model + save_path: Path to save the visualization + """ + ensure_dirs() + model.eval() + + n_classes = NB_LABEL + n_timesteps = 10 + timesteps = np.linspace(1, DIFFU_STEPS, n_timesteps, dtype=int)[::-1] + + fig, axes = plt.subplots(n_classes, n_timesteps, figsize=(2 * n_timesteps, 2 * n_classes)) + + with torch.no_grad(): + # Start from noise + x = torch.randn(n_classes, NB_CHANNEL, IMG_SIZE, IMG_SIZE, device=DEVICE) + + # One sample per class + vec = torch.arange(n_classes, device=DEVICE) + vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(dtype=torch.float32) + + for t in track(range(DIFFU_STEPS, 0, -1), description="Visualizing backward diffusion"): + t_tensor = torch.tensor([[t]] * n_classes, device=DEVICE, dtype=torch.float32) + x = p_xt_1_xt_x0_pred(model, x, t_tensor, vec) + + if t in timesteps: + t_idx = timesteps.tolist().index(t) + for class_idx in range(n_classes): + axes[class_idx, t_idx].imshow(tensor_to_image(x[class_idx]), cmap='gray') + if class_idx == 0: + axes[class_idx, t_idx].set_title(f't={t}') + if t_idx == 0: + axes[class_idx, t_idx].set_ylabel(f'Class {class_idx}') + axes[class_idx, t_idx].set_xticks([]) + axes[class_idx, t_idx].set_yticks([]) + + fig.suptitle('Backward Diffusion Process') + plt.tight_layout() + + if save_path is None: + save_path = os.path.join(PLOTS_DIR, 'backward_diffusion.png') + fig.savefig(save_path) + print(f"Saved backward diffusion visualization to {save_path}") + plt.close(fig) + + +def generate_grid(model: torch.nn.Module, n_per_class: int = 10, save_path: str | None = None): + """ + Generate a grid of samples, organized by class. + + Args: + model: Trained UNet model + n_per_class: Number of samples per class + save_path: Path to save the grid + """ + ensure_dirs() + + # Generate samples + labels = [] + for class_idx in range(NB_LABEL): + labels.extend([class_idx] * n_per_class) + + samples = generate_samples(model, labels=labels) + samples = samples.cpu().numpy() + + # Create grid + fig, axes = plt.subplots(NB_LABEL, n_per_class, figsize=(n_per_class, NB_LABEL)) + + for class_idx in range(NB_LABEL): + for sample_idx in range(n_per_class): + idx = class_idx * n_per_class + sample_idx + axes[class_idx, sample_idx].imshow(samples[idx].transpose(1, 2, 0).squeeze(), cmap='gray') + axes[class_idx, sample_idx].axis('off') + + if sample_idx == 0: + axes[class_idx, sample_idx].set_ylabel(f'{class_idx}') + + fig.suptitle('Generated MNIST Digits') + plt.tight_layout() + + if save_path is None: + save_path = os.path.join(PLOTS_DIR, 'generated_grid.png') + fig.savefig(save_path) + print(f"Saved generated grid to {save_path}") + plt.close(fig) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Sample from trained diffusion model") + parser.add_argument("--checkpoint", type=str, default=os.path.join(CHECKPOINT_DIR, "diffusion_latest.pt"), + help="Path to model checkpoint") + parser.add_argument("--n-samples", type=int, default=10, help="Samples per class") + parser.add_argument("--attention", action="store_true", help="Use attention in model") + parser.add_argument("--forward", action="store_true", help="Visualize forward diffusion") + parser.add_argument("--backward", action="store_true", help="Visualize backward diffusion") + parser.add_argument("--grid", action="store_true", help="Generate sample grid") + parser.add_argument("--all", action="store_true", help="Run all visualizations") + + args = parser.parse_args() + + # Load model + model = UNetMNIST(use_attention=args.attention).to(DEVICE) + + if os.path.exists(args.checkpoint): + model = load_checkpoint(model, os.path.basename(args.checkpoint)) + else: + print(f"Warning: Checkpoint {args.checkpoint} not found. Using untrained model.") + + # Run visualizations + if args.forward or args.all: + visualize_forward_diffusion() + + if args.backward or args.all: + visualize_backward_diffusion(model) + + if args.grid or args.all: + generate_grid(model, n_per_class=args.n_samples) + + # Default: generate grid if no specific option selected + if not (args.forward or args.backward or args.grid or args.all): + generate_grid(model, n_per_class=args.n_samples) diff --git a/src/train_autoencoder.py b/src/train_autoencoder.py new file mode 100644 index 0000000..669ce47 --- /dev/null +++ b/src/train_autoencoder.py @@ -0,0 +1,167 @@ +""" +Training script for MNIST autoencoder. +""" + +from models import Autoencoder, AEModule +from src.utils import ensure_dirs, save_checkpoint +from src.dataloader import get_autoencoder_dataloader +from src.config import DEVICE, BATCH_SIZE, NUM_WORKERS, CHECKPOINT_DIR, PLOTS_DIR +import os +import torch +import torch.nn as nn +import matplotlib.pyplot as plt +from rich.progress import track +import mlflow + +import sys +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +def train_autoencoder( + epochs: int = 10, + learning_rate: float = 1e-3, + batch_size: int = BATCH_SIZE, + latent_channels: int = 1, + checkpoint_path: str | None = None, + run_name: str | None = None, +): + """ + Train the autoencoder model. + + Args: + epochs: Number of training epochs + learning_rate: Learning rate for optimizer + batch_size: Training batch size + latent_channels: Number of channels in latent space + checkpoint_path: Path to checkpoint to resume training from + run_name: Name for this training run (for logging) + """ + ensure_dirs() + + # Initialize model + model = Autoencoder(input_channels=1, latent_channels=latent_channels).to(DEVICE) + + if checkpoint_path: + if not os.path.exists(checkpoint_path): + print(f"Checkpoint not found: {checkpoint_path}") + return + model.load_state_dict(torch.load(checkpoint_path, weights_only=True)) + + if run_name is None: + from datetime import datetime + run_name = f"autoencoder_{datetime.now().strftime('%Y%m%d_%H%M%S')}" + + mlflow.set_experiment("MNIST Autoencoder") + mlflow.start_run(run_name=run_name) + + # Setup training + optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) + # Use MSE loss instead of BCE for better reconstruction of continuous values + criterion = nn.MSELoss() + + # Get dataloader + dataloader = get_autoencoder_dataloader( + batch_size=batch_size, + num_workers=NUM_WORKERS, + ) + + # Training loop + for epoch in range(1, epochs + 1): + model.train() + epoch_loss = 0.0 + + for batch_idx, (x, target) in enumerate(track(dataloader, description=f"Epoch {epoch}/{epochs}")): + optimizer.zero_grad() + + # Forward pass + x_recon = model(x) + loss = criterion(x_recon, target) + + # Backward pass + loss.backward() + + # Gradient clipping to prevent exploding gradients + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + + optimizer.step() + + epoch_loss += loss.item() + + avg_loss = epoch_loss / len(dataloader) + mlflow.log_metric("epoch_loss", avg_loss, step=epoch) + + # Save checkpoint + save_checkpoint(model, f"autoencoder_epoch_{epoch:03d}.pt") + save_checkpoint(model, "autoencoder_latest.pt") + + # Visualize results + visualize_reconstructions(model, dataloader) + + mlflow.end_run() + return model + + +def visualize_reconstructions(model: AEModule, dataloader, n_samples: int = 10): + """Visualize original, latent, and reconstructed images.""" + model.eval() + ensure_dirs() + + # Get a batch of samples + x_batch, _ = next(iter(dataloader)) + x_batch = x_batch[:n_samples] + + with torch.no_grad(): + latent = model.encode(x_batch) + x_recon = model.decode(latent) + + # Create visualization + fig, axes = plt.subplots(3, n_samples + 1, figsize=(2 * n_samples, 6)) + + # Labels + axes[0, 0].text(0.5, 0.5, 'Original', ha='center', va='center', fontsize=12) + axes[0, 0].axis('off') + axes[1, 0].text(0.5, 0.5, 'Latent', ha='center', va='center', fontsize=12) + axes[1, 0].axis('off') + axes[2, 0].text(0.5, 0.5, 'Reconstructed', ha='center', va='center', fontsize=12) + axes[2, 0].axis('off') + + # Plot images + for i in range(n_samples): + axes[0, i + 1].imshow(x_batch[i].cpu().squeeze().numpy(), cmap='gray') + axes[0, i + 1].axis('off') + + axes[1, i + 1].imshow(latent[i].cpu().squeeze().numpy(), cmap='gray') + axes[1, i + 1].axis('off') + + axes[2, i + 1].imshow(x_recon[i].cpu().squeeze().numpy(), cmap='gray') + axes[2, i + 1].axis('off') + + fig.suptitle('Autoencoder Results') + plt.tight_layout() + + save_path = os.path.join(PLOTS_DIR, 'autoencoder_results.png') + fig.savefig(save_path) + plt.close(fig) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Train MNIST autoencoder") + parser.add_argument("--epochs", type=int, default=10, help="Number of epochs") + parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate") + parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="Batch size") + parser.add_argument("--latent-channels", type=int, default=1, help="Latent channels") + parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint") + + args = parser.parse_args() + + torch.multiprocessing.set_start_method("spawn", force=True) + + train_autoencoder( + epochs=args.epochs, + learning_rate=args.lr, + batch_size=args.batch_size, + latent_channels=args.latent_channels, + checkpoint_path=args.checkpoint, + ) diff --git a/src/train_diffusion.py b/src/train_diffusion.py new file mode 100644 index 0000000..cb70650 --- /dev/null +++ b/src/train_diffusion.py @@ -0,0 +1,197 @@ +""" +Training script for MNIST diffusion model. +""" + +from models import UNetMNIST +from src.utils import ( + ensure_dirs, save_checkpoint, tensor_to_image, + figure_to_image, +) +from src.metrics import fid, kl_divergence, jsd +from src.diffusion import p_xt_1_xt_x0_pred +from src.dataloader import get_diffusion_dataloader, get_mnist_dataset +from src.config import ( + DEVICE, EPOCHS, LEARNING_RATE, BATCH_SIZE, NUM_WORKERS, + DIFFU_STEPS, NB_CHANNEL, IMG_SIZE, NB_LABEL, CHECKPOINT_DIR, PLOTS_DIR +) +import os +import torch +import torch.nn as nn +import numpy as np +import matplotlib.pyplot as plt +from rich.progress import track +import mlflow + +import sys +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +def train_diffusion( + epochs: int = EPOCHS, + learning_rate: float = LEARNING_RATE, + batch_size: int = BATCH_SIZE, + predict_x0: bool = True, + use_attention: bool = False, + checkpoint_path: str | None = None, + run_name: str | None = None, +): + """ + Train the diffusion model. + + Args: + epochs: Number of training epochs + learning_rate: Learning rate for optimizer + batch_size: Training batch size + predict_x0: If True, model predicts x0. Otherwise predicts noise. + use_attention: If True, use self-attention in UNet + checkpoint_path: Path to checkpoint to resume training from + run_name: Name for this training run (for logging) + """ + ensure_dirs() + + # Setup run name and logging + if run_name is None: + from datetime import datetime + run_name = f"diffusion_{datetime.now().strftime('%Y%m%d_%H%M%S')}" + + mlflow.set_experiment("MNIST Diffusion") + mlflow.start_run(run_name=run_name) + + # Initialize model + model = UNetMNIST(use_attention=use_attention).to(DEVICE) + + if checkpoint_path and os.path.exists(checkpoint_path): + model.load_state_dict(torch.load(checkpoint_path, weights_only=True)) + print(f"Loaded checkpoint from {checkpoint_path}") + + # Setup training + optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) + criterion = nn.MSELoss() + + # Get dataloader + dataloader = get_diffusion_dataloader( + predict_x0=predict_x0, + batch_size=batch_size, + num_workers=NUM_WORKERS, + ) + + # Training loop + global_step = 0 + for epoch in range(1, epochs + 1): + model.train() + epoch_loss = 0.0 + + for batch_idx, (xt, t, vec, target) in enumerate(track(dataloader, description=f"Epoch {epoch}/{epochs}")): + optimizer.zero_grad() + + # Forward pass + pred = model(xt, t, vec) + loss = criterion(pred, target) + + # Backward pass + loss.backward() + optimizer.step() + + epoch_loss += loss.item() + global_step += 1 + + # Log training loss every 100 steps + if global_step % 100 == 0: + avg_loss = epoch_loss / (batch_idx + 1) + mlflow.log_metric("train_loss", avg_loss, step=global_step) + + avg_loss = epoch_loss / len(dataloader) + mlflow.log_metric("epoch_loss", avg_loss, step=epoch) + + # Save checkpoint every epoch + save_checkpoint(model, f"diffusion_epoch_{epoch:03d}.pt") + save_checkpoint(model, "diffusion_latest.pt") + + # Generate and log sample images + if epoch % 1 == 0: + evaluate_and_log(model, epoch, predict_x0) + + mlflow.end_run() + return model + + +def evaluate_and_log(model: nn.Module, epoch: int, predict_x0: bool = True): + """Generate samples and log metrics.""" + model.eval() + + with torch.no_grad(): + # Generate samples + batch_size = 64 + n_batches = 4 + + fakes = [] + for _ in range(n_batches): + x = torch.randn(batch_size, NB_CHANNEL, IMG_SIZE, IMG_SIZE, device=DEVICE) + vec = torch.randint(0, NB_LABEL, (batch_size,), device=DEVICE) + vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(dtype=torch.float32) + + # Reverse diffusion + for t in range(DIFFU_STEPS, 0, -1): + t_tensor = torch.tensor([[t]] * batch_size, device=DEVICE, dtype=torch.float32) + x = p_xt_1_xt_x0_pred(model, x, t_tensor, vec) + + x = x.cpu().numpy() + x = (x - x.min()) / (x.max() - x.min() + 1e-8) + fakes.append(x) + + fakes = np.concatenate(fakes) + + # Get real samples for comparison + dataset = get_mnist_dataset(train=True) + n_samples = len(fakes) + reals = torch.stack([dataset[i][0] for i in range(n_samples)]).numpy() + reals = reals * 2 - 1 + + # Log metrics + fid_score = fid(reals, fakes) + kl_score = kl_divergence(reals, fakes) + jsd_score = jsd(reals, fakes) + + mlflow.log_metric("FID", fid_score, step=epoch) + mlflow.log_metric("KL Divergence", kl_score, step=epoch) + mlflow.log_metric("JSD", jsd_score, step=epoch) + + # Log sample images + fig, axes = plt.subplots(4, 8, figsize=(16, 8)) + for i, ax in enumerate(axes.flat): + if i < len(fakes): + ax.imshow(fakes[i].transpose(1, 2, 0).squeeze(), cmap='gray') + ax.axis('off') + fig.suptitle(f"Generated Samples - Epoch {epoch}") + plt.tight_layout() + + mlflow.log_figure(fig, f"samples_epoch_{epoch:03d}.png") + + # Save to plots folder + fig.savefig(os.path.join(PLOTS_DIR, f"samples_epoch_{epoch:03d}.png")) + plt.close(fig) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Train MNIST diffusion model") + parser.add_argument("--epochs", type=int, default=EPOCHS, help="Number of epochs") + parser.add_argument("--lr", type=float, default=LEARNING_RATE, help="Learning rate") + parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="Batch size") + parser.add_argument("--attention", action="store_true", help="Use self-attention") + parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint") + parser.add_argument("--name", type=str, default=None, help="Run name") + + args = parser.parse_args() + + torch.multiprocessing.set_start_method("spawn", force=True) + + train_diffusion( + epochs=args.epochs, + learning_rate=args.lr, + batch_size=args.batch_size, + use_attention=args.attention, + checkpoint_path=args.checkpoint, + run_name=args.name, + ) diff --git a/src/utils.py b/src/utils.py new file mode 100644 index 0000000..f727a68 --- /dev/null +++ b/src/utils.py @@ -0,0 +1,110 @@ +""" +Utility functions for MNIST diffusion. +""" + +import os +import io +import numpy as np +import torch +import scipy.linalg +from PIL import Image +import matplotlib.pyplot as plt + +from .config import CHECKPOINT_DIR, PLOTS_DIR + + +def ensure_dirs(): + """Create necessary directories if they don't exist.""" + os.makedirs(CHECKPOINT_DIR, exist_ok=True) + os.makedirs(PLOTS_DIR, exist_ok=True) + + +def tensor_to_image(tensor: torch.Tensor) -> np.ndarray: + """ + Convert a tensor to a numpy image array. + + Args: + tensor: Image tensor [C, H, W] + + Returns: + Numpy array [H, W, C] normalized to [0, 1] + """ + img = tensor.clone().detach().cpu().numpy().transpose(1, 2, 0) + img -= img.min() + img /= img.max() + 1e-8 + return img + + +def tensor_to_images(tensor: torch.Tensor) -> np.ndarray: + """ + Convert a batch of tensors to numpy image arrays. + + Args: + tensor: Batch of image tensors [B, C, H, W] + + Returns: + Numpy array [B, H, W, C] normalized to [0, 1] + """ + img = tensor.clone().detach().cpu().numpy().transpose(0, 2, 3, 1) + img -= np.min(img, axis=(1, 2, 3), keepdims=True) + img /= np.max(img, axis=(1, 2, 3), keepdims=True) + 1e-8 + return img + + +def figure_to_image(figure: plt.Figure) -> np.ndarray: + """ + Convert a matplotlib figure to a numpy image array. + + Args: + figure: Matplotlib figure + + Returns: + Numpy array of the figure image + """ + buf = io.BytesIO() + figure.savefig(buf, format='png') + buf.seek(0) + image = np.array(Image.open(buf)) + return image + + +def save_checkpoint(model: torch.nn.Module, filename: str): + """ + Save model checkpoint. + + Args: + model: Model to save + filename: Filename (will be saved in CHECKPOINT_DIR) + """ + ensure_dirs() + path = os.path.join(CHECKPOINT_DIR, filename) + torch.save(model.state_dict(), path) + + +def load_checkpoint(model: torch.nn.Module, filename: str) -> torch.nn.Module: + """ + Load model checkpoint. + + Args: + model: Model architecture to load weights into + filename: Filename (loaded from CHECKPOINT_DIR) + + Returns: + Model with loaded weights + """ + path = os.path.join(CHECKPOINT_DIR, filename) + model.load_state_dict(torch.load(path, weights_only=True)) + return model + + +def save_plot(figure: plt.Figure, filename: str): + """ + Save a matplotlib figure to the plots directory. + + Args: + figure: Matplotlib figure to save + filename: Filename (will be saved in PLOTS_DIR) + """ + ensure_dirs() + path = os.path.join(PLOTS_DIR, filename) + figure.savefig(path) |