train_diffusion.py

import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt
from diffusers import DDPMScheduler
from distributions import LorenzDataset, RingDataset, TwoPeaksDataset, TwoMoonsDataset

## Model


class Diffusion(nn.Module):
    def __init__(self, num_steps=100, dim=3, hidden_dim=64, prediction_type="v_prediction"):
        super(Diffusion, self).__init__()
        self.nonlinear = nn.ReLU()
        self.num_steps = num_steps
        self.dim = dim
        prediction_type = self._to_prediction_type(prediction_type)
        if prediction_type not in {"sample", "epsilon", "v_prediction"}:
            raise ValueError(f"Unsupported prediction_type: {prediction_type}")
        self.prediction_type = prediction_type
        self.mlp = nn.Sequential(
            nn.Linear(dim + 1, hidden_dim),
            self.nonlinear,
            nn.Linear(hidden_dim, hidden_dim),
            self.nonlinear,
            nn.Linear(hidden_dim, dim),
        )

        # clip_sample=False avoids forced truncation to [-1, 1].
        self.scheduler = DDPMScheduler(
            num_train_timesteps=num_steps,
            beta_schedule="squaredcos_cap_v2",
            prediction_type=prediction_type,
            clip_sample=False,
        )
        self.register_buffer(
            "alphas_cumprod",
            self.scheduler.alphas_cumprod.to(torch.float32),
            persistent=False,
        )

    def x_predictor(self, x_t, t):
        model_input = torch.cat([x_t, t * 2 - 1], dim=-1)
        return self.mlp(model_input) + x_t

    @staticmethod
    def _to_prediction_type(prediction_type):
        if prediction_type == "x_prediction":
            return "sample"
        return prediction_type

    def _timesteps_from_t(self, t):
        return (t * (self.num_steps - 1)).long().clamp(0, self.num_steps - 1)

    def _eps_from_x0(self, x_t, x0_pred, alpha_t):
        return (x_t - alpha_t.sqrt() * x0_pred) / (1 - alpha_t).sqrt().clamp_min(1e-8)

    def _model_output_from_x0(self, x_t, x0_pred, alpha_t):
        eps_pred = self._eps_from_x0(x_t, x0_pred, alpha_t)
        if self.prediction_type == "sample":
            return x0_pred
        if self.prediction_type == "epsilon":
            return eps_pred
        if self.prediction_type == "v_prediction":
            return alpha_t.sqrt() * eps_pred - (1 - alpha_t).sqrt() * x0_pred
        raise ValueError(f"Unsupported prediction_type: {self.prediction_type}")

    def _x0_from_model_output(self, x_t, model_output, alpha_t):
        if self.prediction_type == "sample":
            return model_output
        if self.prediction_type == "epsilon":
            return (x_t - (1 - alpha_t).sqrt() * model_output) / alpha_t.sqrt().clamp_min(
                1e-8
            )
        if self.prediction_type == "v_prediction":
            return alpha_t.sqrt() * x_t - (1 - alpha_t).sqrt() * model_output
        raise ValueError(f"Unsupported prediction_type: {self.prediction_type}")

    def _target_from_clean(self, x0, eps, alpha_t):
        if self.prediction_type == "sample":
            return x0
        if self.prediction_type == "epsilon":
            return eps
        if self.prediction_type == "v_prediction":
            return alpha_t.sqrt() * eps - (1 - alpha_t).sqrt() * x0
        raise ValueError(f"Unsupported prediction_type: {self.prediction_type}")

    def model_output(self, x_t, t):
        alpha_t = self.alpha_t(t)
        x0_pred = self.x_predictor(x_t, t.unsqueeze(-1))
        return self._model_output_from_x0(x_t, x0_pred, alpha_t)

    def score(self, x, t: float = 0.1):
        if isinstance(t, float):
            batch_size = x.shape[0]
            t = torch.ones(batch_size, device=x.device) * t
        alpha_t = self.alpha_t(t)
        model_output = self.model_output(x, t)
        x0_pred = self._x0_from_model_output(x, model_output, alpha_t)
        eps_pred = self._eps_from_x0(x, x0_pred, alpha_t)
        return -eps_pred / (1 - alpha_t).sqrt()

    def forward(self, x, t):
        x_t, eps = self.diffuse(x, t)
        alpha_t = self.alpha_t(t)
        pred = self.model_output(x_t, t)
        target = self._target_from_clean(x, eps, alpha_t)
        return target, pred

    def alpha_t(self, t):
        t = self._timesteps_from_t(t)
        alpha_t = self.alphas_cumprod[t].unsqueeze(-1).to(t.device)
        return alpha_t

    def diffuse(self, x, t):
        alpha_t = self.alpha_t(t)
        eps = torch.randn_like(x)
        x_t = alpha_t.sqrt() * x + (1 - alpha_t).sqrt() * eps
        return x_t, eps

    def predict_x0(self, x, t: float):
        with torch.no_grad():
            alpha_t = self.alpha_t(t)
            model_output = self.model_output(x, t)
            return self._x0_from_model_output(x, model_output, alpha_t)

    def sample(self, num_sample, device=None):
        if device is None:
            device = next(self.parameters()).device
        x = torch.randn(num_sample, self.dim, device=device)
        self.scheduler.set_timesteps(self.num_steps, device=device)
        for step_t in self.scheduler.timesteps:
            t_norm = (
                torch.full((num_sample,), step_t.item(), device=device)
                / (self.num_steps - 1)
            )
            with torch.no_grad():
                model_output = self.model_output(x, t_norm)
            x = self.scheduler.step(model_output, step_t, x).prev_sample
        return x


def train(
    dataset,
    num_epochs=500,
    device="none",
    batch_size=2048,
    lr=1e-2,
    prediction_type="v_prediction",
):
    from tqdm import tqdm

    if device == "none":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    model = Diffusion(dim=dataset.dim, prediction_type=prediction_type).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5)

    losses = []
    lf = nn.MSELoss()

    for epoch in tqdm(range(num_epochs)):
        acc_loss = 0
        for x in dataloader:
            optimizer.zero_grad()
            x = x.to(device)
            t = torch.rand(x.shape[0]).to(device)
            target, pred = model.forward(x, t)
            loss = lf(target, pred)
            loss.backward()
            optimizer.step()
            acc_loss += loss.item()
        losses.append(acc_loss / len(dataloader))

    model.eval()

    return model, losses


def _ensure_dir(path):
    os.makedirs(path, exist_ok=True)


def test_sample(model, dataset, args, output_dir, num_sample=100):
    model.eval()
    sample_device = next(model.parameters()).device
    sampled = model.sample(num_sample, device=sample_device).cpu()
    if args.model == "lorenz":
        marker = "-"
    else:
        marker = "."
    plt.plot(
        dataset.time_series[:, 0],
        dataset.time_series[:, 1],
        marker,
        alpha=0.5,
        label="True",
    )
    plt.plot(sampled[:, 0], sampled[:, 1], ".", alpha=0.5, label="Sampled")
    plt.title("Sampled")
    plt.xlabel("x")
    plt.ylabel("y")
    plt.legend()
    plt.savefig(os.path.join(output_dir, "diffusion_sampled.png"))
    plt.close()


def plot_losses(losses, output_dir):
    plt.plot(losses)
    plt.title("Losses")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.savefig(os.path.join(output_dir, "diffusion_loss.png"))
    plt.close()


def save_model(model, output_dir):
    save_path = os.path.join(output_dir, "diffusion_model.pth")
    payload = {"state_dict": model.state_dict(), "prediction_type": model.prediction_type}
    torch.save(payload, save_path)
    print(f"Model saved to {save_path}")


if __name__ == "__main__":
    # set seed for reproducibility
    np.random.seed(0)
    torch.manual_seed(0)
    import argparse

    # set model: choose from lorenz, ring, two_peaks
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, default="lorenz")
    parser.add_argument("--num_epochs", type=int, default=512)
    parser.add_argument("--num_sample", type=int, default=1000)
    parser.add_argument("--device", type=str, default="none")
    parser.add_argument("--batch_size", type=int, default=2048)
    parser.add_argument("--lr", type=float, default=1e-2)
    parser.add_argument("--output_dir", type=str, default=None)
    parser.add_argument(
        "--prediction_type",
        type=str,
        default="v_prediction",
        choices=["sample", "x_prediction", "epsilon", "v_prediction"],
    )
    args = parser.parse_args()
    if args.model == "lorenz":
        dataset = LorenzDataset()
    elif args.model == "ring":
        dataset = RingDataset()
    elif args.model == "two_peaks":
        dataset = TwoPeaksDataset()
    elif args.model == "two_moons":
        dataset = TwoMoonsDataset()
    else:
        raise ValueError(f"Model {args.model} not supported")

    output_dir = args.output_dir or f"./results/{args.model}"
    _ensure_dir(output_dir)
    prediction_type = Diffusion._to_prediction_type(args.prediction_type)
    model, losses = train(
        dataset,
        num_epochs=args.num_epochs,
        device=args.device,
        batch_size=args.batch_size,
        lr=args.lr,
        prediction_type=prediction_type,
    )
    plot_losses(losses, output_dir)
    test_sample(model, dataset, args, output_dir, num_sample=args.num_sample)
    save_model(model, output_dir)