path: root/src/train_diffusion.py

"""
Training script for MNIST diffusion model.
"""

from models import UNetMNIST
from src.utils import ensure_dirs, save_checkpoint
from src.metrics import fid, kl_divergence, jsd
from src.diffusion import p_xt_1_xt_x0_pred
from src.dataloader import get_diffusion_dataloaders, get_mnist_dataset
from src.config import (
    DEVICE, EPOCHS, LEARNING_RATE, BATCH_SIZE, NUM_WORKERS,
    DIFFU_STEPS, NB_CHANNEL, IMG_SIZE, NB_LABEL
)
import os
import torch
import torch.nn as nn
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
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 | None = None,
    learning_rate: float = LEARNING_RATE,
    batch_size: int = BATCH_SIZE,
    predict_x0: bool = True,
    use_attention: bool = False,
    val_split: float = 0.1,
    test_split: float = 0.1,
    patience: int | None = 5,
    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
        val_split: Fraction of data for validation
        test_split: Fraction of data for testing
        patience: Number of epochs to wait for validation loss improvement before stopping
        checkpoint_path: Path to checkpoint to resume training from
        run_name: Name for this training run (for logging)
    """
    if epochs is None and patience is None:
        raise ValueError("Must specify either epochs or patience for training")
    if patience is not None and patience <= 0:
        raise ValueError("Patience must be a positive integer")

    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 dataloaders
    train_loader, val_loader, test_loader = get_diffusion_dataloaders(
        predict_x0=predict_x0,
        batch_size=batch_size,
        num_workers=NUM_WORKERS,
        val_split=val_split,
        test_split=test_split,
    )

    # Early stopping tracking
    best_val_loss = float('inf')
    epochs_without_improvement = 0

    # Training loop
    global_step = 0
    epoch = 0
    while epochs is None or epoch < epochs:
        epoch += 1

        model.train()
        epoch_loss = 0.0

        if epochs:
            description = f"Epoch {epoch}/{epochs}"
        else:
            description = f"Epoch {epoch}"
        if patience:
            description += f" (Patience: {patience-epochs_without_improvement})"

        for batch_idx, (xt, t, vec, target) in enumerate(track(train_loader, description=description)):
            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_step", avg_loss, step=global_step)

        avg_loss = epoch_loss / len(train_loader)
        mlflow.log_metric("train/loss", avg_loss, step=epoch)

        val_loss = evaluate_diffusion_loss(model, val_loader, criterion)
        test_loss = evaluate_diffusion_loss(model, test_loader, criterion)

        mlflow.log_metric("val/loss", val_loss, step=epoch)
        mlflow.log_metric("test/loss", test_loss, step=epoch)

        # Early stopping check
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            epochs_without_improvement = 0
            # Save best checkpoint
            save_checkpoint(model, "diffusion_best.pt")
        else:
            epochs_without_improvement += 1

        # Save regular checkpoints
        save_checkpoint(model, f"diffusion_epoch_{epoch:03d}.pt")
        save_checkpoint(model, "diffusion_latest.pt")

        # Stop if no improvement
        if patience and epochs_without_improvement >= patience:
            print(f"Early stopping: No improvement for {patience} epochs")
            break

        # Generate and log sample images and metrics on test split
        if epoch % 1 == 0:
            evaluate_and_log(model, epoch, test_loader, predict_x0)

    mlflow.end_run()
    return model


def evaluate_diffusion_loss(
    model: nn.Module,
    dataloader,
    criterion,
) -> float:
    """Evaluate diffusion model and return average loss."""
    model.eval()
    total_loss = 0.0

    with torch.no_grad():
        for xt, t, vec, target in dataloader:
            pred = model(xt, t, vec)
            loss = criterion(pred, target)
            total_loss += loss.item()

    return total_loss / len(dataloader)


def evaluate_and_log(
    model: nn.Module,
    epoch: int,
    test_loader,
    predict_x0: bool = True,
):
    """Generate samples and log metrics evaluated on test split."""
    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 from test split
        n_samples = len(fakes)
        base_dataset = getattr(test_loader.dataset, "dataset", None)
        if base_dataset is None:
            base_dataset = get_mnist_dataset(train=True)
        n_samples = min(n_samples, len(base_dataset))
        reals = torch.stack([base_dataset[i][0] for i in range(n_samples)]).cpu().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("test/FID", fid_score, step=epoch)
        mlflow.log_metric("test/KL Divergence", kl_score, step=epoch)
        mlflow.log_metric("test/JSD", jsd_score, step=epoch)

        # Log sample images using plotly
        fig = make_subplots(rows=4, cols=8, horizontal_spacing=0.01, vertical_spacing=0.02)
        for i in range(min(32, len(fakes))):
            row = i // 8 + 1
            col = i % 8 + 1
            img = fakes[i].transpose(1, 2, 0).squeeze()[::-1]
            fig.add_trace(
                go.Heatmap(z=img, colorscale='gray', showscale=False),
                row=row, col=col
            )
        fig.update_layout(
            title_text=f"Generated Samples - Epoch {epoch}",
            width=800,
            height=400,
            showlegend=False
        )
        fig.update_xaxes(showticklabels=False, showgrid=False, zeroline=False)
        fig.update_yaxes(showticklabels=False, showgrid=False, zeroline=False)

        mlflow.log_figure(fig, f"epoch_{epoch:03d}.png")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Train MNIST diffusion model")
    parser.add_argument("--epochs", type=int, default=None, 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("--val-split", type=float, default=0.1, help="Validation split fraction")
    parser.add_argument("--test-split", type=float, default=0.1, help="Test split fraction")
    parser.add_argument("--patience", type=int, default=5, help="Early stopping patience")
    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,
        val_split=args.val_split,
        test_split=args.test_split,
        patience=args.patience,
        checkpoint_path=args.checkpoint,
        run_name=args.name,
    )