Johnny Fernandes
781 lines
31 KiB
Python
781 lines
31 KiB
Python
import os
|
||
import time
|
||
from pathlib import Path
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from torch.utils.data import DataLoader
|
||
from torchvision.utils import save_image
|
||
from tqdm import tqdm
|
||
|
||
from src.training.ema import EMA
|
||
from src.training.fid import FIDEvaluator
|
||
|
||
if hasattr(torch.amp, "GradScaler"):
|
||
_GradScaler = torch.amp.GradScaler
|
||
_autocast = torch.amp.autocast
|
||
else:
|
||
from torch.cuda.amp import GradScaler as _GS, autocast as _AC
|
||
_GradScaler = lambda device="", enabled=True, **kw: _GS(**kw)
|
||
_autocast = lambda device_type="", enabled=True, **kw: _AC(enabled=enabled, **kw)
|
||
|
||
|
||
def _save_samples(generator_ema, samples_dir: Path, epoch: int, *, fixed_noise: torch.Tensor, device) -> None:
|
||
samples_dir.mkdir(parents=True, exist_ok=True)
|
||
with torch.no_grad():
|
||
imgs = generator_ema.model(fixed_noise.to(device)) # EMA model, [-1, 1]
|
||
imgs = (imgs.clamp(-1, 1) + 1.0) / 2.0 # -> [0, 1]
|
||
save_image(imgs, samples_dir / f"epoch_{epoch:04d}.png", nrow=4)
|
||
|
||
|
||
def train_dcgan(
|
||
generator,
|
||
discriminator,
|
||
train_dataset,
|
||
cfg: dict,
|
||
*,
|
||
save_dir,
|
||
run_name: str,
|
||
device: str = "cuda",
|
||
) -> dict:
|
||
"""Vanilla DCGAN training loop with BCE loss (Radford et al., 2015).
|
||
|
||
Used as the Phase 1 baseline for cheap pipeline ablations. No gradient
|
||
penalty, no n_critic, single G/D step per batch.
|
||
"""
|
||
device = torch.device(device if torch.cuda.is_available() else "cpu")
|
||
generator = generator.to(device)
|
||
discriminator = discriminator.to(device)
|
||
|
||
n_g = sum(p.numel() for p in generator.parameters() if p.requires_grad)
|
||
n_d = sum(p.numel() for p in discriminator.parameters() if p.requires_grad)
|
||
print(f"Generator: {n_g:,} params Discriminator: {n_d:,} params")
|
||
|
||
epochs = cfg["epochs"]
|
||
batch_size = cfg["batch_size"]
|
||
lr_g = cfg.get("lr_g", 2e-4)
|
||
lr_d = cfg.get("lr_d", 2e-4)
|
||
beta1 = cfg.get("beta1", 0.5)
|
||
beta2 = cfg.get("beta2", 0.999)
|
||
latent_dim = cfg.get("latent_dim", 100)
|
||
ema_decay = cfg.get("ema_decay", 0.9999)
|
||
sample_interval = cfg.get("sample_interval", 10)
|
||
fid_interval = cfg.get("fid_interval", 25)
|
||
fid_n_real = cfg.get("fid_n_real", 5000)
|
||
|
||
loader = DataLoader(
|
||
train_dataset, batch_size=batch_size, shuffle=True,
|
||
num_workers=cfg.get("num_workers", min(4, os.cpu_count() or 1)),
|
||
pin_memory=(device.type == "cuda"), drop_last=True,
|
||
)
|
||
opt_g = torch.optim.Adam(generator.parameters(), lr=lr_g, betas=(beta1, beta2))
|
||
opt_d = torch.optim.Adam(discriminator.parameters(), lr=lr_d, betas=(beta1, beta2))
|
||
bce = nn.BCEWithLogitsLoss()
|
||
|
||
use_amp = device.type == "cuda"
|
||
scaler_g = _GradScaler("cuda", enabled=use_amp)
|
||
scaler_d = _GradScaler("cuda", enabled=use_amp)
|
||
|
||
ema = EMA(generator, decay=ema_decay)
|
||
|
||
# Fixed noise for consistent sample tracking across epochs
|
||
fixed_noise = torch.randn(16, latent_dim, 1, 1, device=device)
|
||
|
||
save_dir = Path(save_dir)
|
||
save_dir.mkdir(parents=True, exist_ok=True)
|
||
samples_dir = save_dir.parent / "samples" / run_name
|
||
|
||
fid_eval = FIDEvaluator(train_dataset, n_real=fid_n_real, device=str(device),
|
||
num_workers=cfg.get("num_workers", min(4, os.cpu_count() or 1)))
|
||
|
||
history = {"g_loss": [], "d_loss": [], "d_real": [], "d_fake": [], "fid": {}}
|
||
best_fid = float("inf")
|
||
print(f"Device: {device} AMP: {use_amp} Batches/epoch: {len(loader)}")
|
||
|
||
# Linear LR decay from epoch epochs//2 to epochs
|
||
decay_start = epochs // 2
|
||
sched_g = torch.optim.lr_scheduler.LambdaLR(
|
||
opt_g, lr_lambda=lambda ep: max(0.0, 1.0 - max(ep - decay_start, 0) / (epochs - decay_start)))
|
||
sched_d = torch.optim.lr_scheduler.LambdaLR(
|
||
opt_d, lr_lambda=lambda ep: max(0.0, 1.0 - max(ep - decay_start, 0) / (epochs - decay_start)))
|
||
|
||
t_start = time.time()
|
||
|
||
for epoch in range(1, epochs + 1):
|
||
generator.train()
|
||
discriminator.train()
|
||
g_sum = d_sum = real_sum = fake_sum = 0.0
|
||
n_batches = 0
|
||
|
||
for imgs in tqdm(loader, desc=f"Epoch {epoch}/{epochs}", leave=False):
|
||
imgs = imgs.to(device)
|
||
bsz = imgs.size(0)
|
||
real_labels = torch.ones(bsz, device=device)
|
||
fake_labels = torch.zeros(bsz, device=device)
|
||
|
||
# ── Discriminator step ────────────────────────────────────────
|
||
noise = torch.randn(bsz, latent_dim, 1, 1, device=device)
|
||
with _autocast("cuda", enabled=use_amp):
|
||
fake = generator(noise).detach()
|
||
d_real = discriminator(imgs)
|
||
d_fake = discriminator(fake)
|
||
d_loss = bce(d_real, real_labels) + bce(d_fake, fake_labels)
|
||
opt_d.zero_grad()
|
||
scaler_d.scale(d_loss).backward()
|
||
scaler_d.step(opt_d)
|
||
scaler_d.update()
|
||
|
||
# ── Generator step ────────────────────────────────────────────
|
||
noise = torch.randn(bsz, latent_dim, 1, 1, device=device)
|
||
with _autocast("cuda", enabled=use_amp):
|
||
g_loss = bce(discriminator(generator(noise)), real_labels)
|
||
opt_g.zero_grad()
|
||
scaler_g.scale(g_loss).backward()
|
||
scaler_g.step(opt_g)
|
||
scaler_g.update()
|
||
ema.update(generator)
|
||
|
||
g_sum += g_loss.item()
|
||
d_sum += d_loss.item()
|
||
real_sum += d_real.mean().item()
|
||
fake_sum += d_fake.mean().item()
|
||
n_batches += 1
|
||
|
||
avg_g = g_sum / n_batches
|
||
avg_d = d_sum / n_batches
|
||
avg_r = real_sum / n_batches
|
||
avg_f = fake_sum / n_batches
|
||
history["g_loss"].append(avg_g)
|
||
history["d_loss"].append(avg_d)
|
||
history["d_real"].append(avg_r)
|
||
history["d_fake"].append(avg_f)
|
||
print(
|
||
f"[{epoch:03d}/{epochs}] "
|
||
f"G: {avg_g:.4f} D: {avg_d:.4f} D(real): {avg_r:.4f} D(fake): {avg_f:.4f}"
|
||
)
|
||
|
||
if epoch % sample_interval == 0:
|
||
_save_samples(ema, samples_dir, epoch, fixed_noise=fixed_noise, device=device)
|
||
|
||
if epoch % fid_interval == 0:
|
||
ema.model.eval()
|
||
with torch.no_grad():
|
||
fake_imgs = torch.cat([
|
||
ema.model(torch.randn(64, latent_dim, 1, 1, device=device))
|
||
for _ in range(fid_n_real // 64 + 1)
|
||
])[:fid_n_real]
|
||
fid_score = fid_eval.compute(fake_imgs)
|
||
history["fid"][epoch] = fid_score
|
||
print(f" FID @ epoch {epoch}: {fid_score:.2f}")
|
||
|
||
if fid_score < best_fid:
|
||
best_fid = fid_score
|
||
torch.save(generator.state_dict(), save_dir / f"{run_name}_best_g.pt")
|
||
torch.save(ema.model.state_dict(), save_dir / f"{run_name}_best_ema.pt")
|
||
|
||
sched_g.step()
|
||
sched_d.step()
|
||
|
||
torch.save(generator.state_dict(), save_dir / f"{run_name}_final_g.pt")
|
||
torch.save(discriminator.state_dict(), save_dir / f"{run_name}_final_d.pt")
|
||
torch.save(ema.model.state_dict(), save_dir / f"{run_name}_final_ema.pt")
|
||
history["train_time_s"] = time.time() - t_start
|
||
return history
|
||
|
||
|
||
def _gradient_penalty(critic, real: torch.Tensor, fake: torch.Tensor, device) -> torch.Tensor:
|
||
"""Two-sided gradient penalty (Gulrajani et al., 2017)."""
|
||
bsz = real.size(0)
|
||
eps = torch.rand(bsz, 1, 1, 1, device=device)
|
||
interp = (eps * real + (1.0 - eps) * fake).requires_grad_(True)
|
||
d_interp = critic(interp)
|
||
grad = torch.autograd.grad(
|
||
outputs=d_interp,
|
||
inputs=interp,
|
||
grad_outputs=torch.ones_like(d_interp),
|
||
create_graph=True,
|
||
retain_graph=True,
|
||
)[0]
|
||
return ((grad.norm(2, dim=[1, 2, 3]) - 1) ** 2).mean()
|
||
|
||
|
||
def train_wgan(
|
||
generator,
|
||
critic,
|
||
train_dataset,
|
||
cfg: dict,
|
||
*,
|
||
save_dir,
|
||
run_name: str,
|
||
device: str = "cuda",
|
||
) -> dict:
|
||
"""WGAN-GP training loop (Gulrajani et al., 2017).
|
||
|
||
Used for Phase 2.2–2.4. Gradient penalty replaces weight clipping.
|
||
The critic runs in float32 to keep GP gradient computation numerically
|
||
stable; AMP is used only for the generator forward/backward.
|
||
"""
|
||
device = torch.device(device if torch.cuda.is_available() else "cpu")
|
||
generator = generator.to(device)
|
||
critic = critic.to(device)
|
||
|
||
n_g = sum(p.numel() for p in generator.parameters() if p.requires_grad)
|
||
n_c = sum(p.numel() for p in critic.parameters() if p.requires_grad)
|
||
print(f"Generator: {n_g:,} params Critic: {n_c:,} params")
|
||
|
||
epochs = cfg["epochs"]
|
||
batch_size = cfg["batch_size"]
|
||
lr_g = cfg.get("lr_g", 1e-4)
|
||
lr_d = cfg.get("lr_d", 1e-4)
|
||
beta1 = cfg.get("beta1", 0.0)
|
||
beta2 = cfg.get("beta2", 0.9)
|
||
latent_dim = cfg.get("latent_dim", 128)
|
||
n_critic = cfg.get("n_critic", 5)
|
||
gp_lambda = cfg.get("gp_lambda", 10)
|
||
ema_decay = cfg.get("ema_decay", 0.9999)
|
||
sample_interval = cfg.get("sample_interval", 10)
|
||
fid_interval = cfg.get("fid_interval", 25)
|
||
fid_n_real = cfg.get("fid_n_real", 5000)
|
||
|
||
loader = DataLoader(
|
||
train_dataset, batch_size=batch_size, shuffle=True,
|
||
num_workers=cfg.get("num_workers", min(4, os.cpu_count() or 1)),
|
||
pin_memory=(device.type == "cuda"), drop_last=True,
|
||
)
|
||
opt_g = torch.optim.Adam(generator.parameters(), lr=lr_g, betas=(beta1, beta2))
|
||
opt_c = torch.optim.Adam(critic.parameters(), lr=lr_d, betas=(beta1, beta2))
|
||
|
||
use_amp = device.type == "cuda"
|
||
scaler_g = _GradScaler("cuda", enabled=use_amp)
|
||
|
||
ema = EMA(generator, decay=ema_decay)
|
||
|
||
# Fixed noise for consistent sample tracking across epochs
|
||
fixed_noise = torch.randn(16, latent_dim, 1, 1, device=device)
|
||
|
||
save_dir = Path(save_dir)
|
||
save_dir.mkdir(parents=True, exist_ok=True)
|
||
samples_dir = save_dir.parent / "samples" / run_name
|
||
|
||
fid_eval = FIDEvaluator(train_dataset, n_real=fid_n_real, device=str(device),
|
||
num_workers=cfg.get("num_workers", min(4, os.cpu_count() or 1)))
|
||
|
||
history = {"g_loss": [], "w_dist": [], "gp": [], "d_real": [], "d_fake": [], "fid": {}}
|
||
best_fid = float("inf")
|
||
print(f"Device: {device} AMP (G only): {use_amp} Batches/epoch: {len(loader)} n_critic: {n_critic}")
|
||
|
||
# Linear LR decay from epoch epochs//2 to epochs
|
||
decay_start = epochs // 2
|
||
sched_g = torch.optim.lr_scheduler.LambdaLR(
|
||
opt_g, lr_lambda=lambda ep: max(0.0, 1.0 - max(ep - decay_start, 0) / (epochs - decay_start)))
|
||
sched_c = torch.optim.lr_scheduler.LambdaLR(
|
||
opt_c, lr_lambda=lambda ep: max(0.0, 1.0 - max(ep - decay_start, 0) / (epochs - decay_start)))
|
||
|
||
t_start = time.time()
|
||
|
||
for epoch in range(1, epochs + 1):
|
||
generator.train()
|
||
critic.train()
|
||
g_sum = w_sum = gp_sum = real_sum = fake_sum = 0.0
|
||
n_c_steps = n_g_steps = 0
|
||
|
||
for batch_idx, real in enumerate(tqdm(loader, desc=f"Epoch {epoch}/{epochs}", leave=False)):
|
||
real = real.to(device)
|
||
bsz = real.size(0)
|
||
|
||
# ── Critic step (every batch) ─────────────────────────────────
|
||
# Run critic in float32 — GP requires double-precision gradients
|
||
# and AMP can degrade stability here.
|
||
opt_c.zero_grad()
|
||
with torch.no_grad():
|
||
fake = generator(torch.randn(bsz, latent_dim, 1, 1, device=device))
|
||
|
||
real_f32 = real.float()
|
||
fake_f32 = fake.float().detach()
|
||
|
||
d_real = critic(real_f32)
|
||
d_fake = critic(fake_f32)
|
||
gp = _gradient_penalty(critic, real_f32, fake_f32.detach(), device)
|
||
c_loss = d_fake.mean() - d_real.mean() + gp_lambda * gp
|
||
c_loss.backward()
|
||
opt_c.step()
|
||
|
||
w_dist = (d_real.mean() - d_fake.mean()).item()
|
||
w_sum += w_dist
|
||
gp_sum += gp.item()
|
||
real_sum += d_real.mean().item()
|
||
fake_sum += d_fake.mean().item()
|
||
n_c_steps += 1
|
||
|
||
# ── Generator step (every n_critic batches) ───────────────────
|
||
if (batch_idx + 1) % n_critic == 0:
|
||
opt_g.zero_grad()
|
||
with _autocast("cuda", enabled=use_amp):
|
||
fake = generator(torch.randn(bsz, latent_dim, 1, 1, device=device))
|
||
g_loss = -critic(fake.float()).mean()
|
||
scaler_g.scale(g_loss).backward()
|
||
scaler_g.step(opt_g)
|
||
scaler_g.update()
|
||
ema.update(generator)
|
||
g_sum += g_loss.item()
|
||
n_g_steps += 1
|
||
|
||
avg_w = w_sum / max(n_c_steps, 1)
|
||
avg_gp = gp_sum / max(n_c_steps, 1)
|
||
avg_g = g_sum / max(n_g_steps, 1)
|
||
avg_r = real_sum / max(n_c_steps, 1)
|
||
avg_f = fake_sum / max(n_c_steps, 1)
|
||
history["g_loss"].append(avg_g)
|
||
history["w_dist"].append(avg_w)
|
||
history["gp"].append(avg_gp)
|
||
history["d_real"].append(avg_r)
|
||
history["d_fake"].append(avg_f)
|
||
print(
|
||
f"[{epoch:03d}/{epochs}] "
|
||
f"G: {avg_g:.4f} W-dist: {avg_w:.4f} GP: {avg_gp:.4f} "
|
||
f"C(real): {avg_r:.4f} C(fake): {avg_f:.4f}"
|
||
)
|
||
|
||
if epoch % sample_interval == 0:
|
||
_save_samples(ema, samples_dir, epoch, fixed_noise=fixed_noise, device=device)
|
||
|
||
if epoch % fid_interval == 0:
|
||
ema.model.eval()
|
||
with torch.no_grad():
|
||
fake_imgs = torch.cat([
|
||
ema.model(torch.randn(64, latent_dim, 1, 1, device=device))
|
||
for _ in range(fid_n_real // 64 + 1)
|
||
])[:fid_n_real]
|
||
fid_score = fid_eval.compute(fake_imgs)
|
||
history["fid"][epoch] = fid_score
|
||
print(f" FID @ epoch {epoch}: {fid_score:.2f}")
|
||
|
||
if fid_score < best_fid:
|
||
best_fid = fid_score
|
||
torch.save(generator.state_dict(), save_dir / f"{run_name}_best_g.pt")
|
||
torch.save(ema.model.state_dict(), save_dir / f"{run_name}_best_ema.pt")
|
||
|
||
sched_g.step()
|
||
sched_c.step()
|
||
|
||
torch.save(generator.state_dict(), save_dir / f"{run_name}_final_g.pt")
|
||
torch.save(critic.state_dict(), save_dir / f"{run_name}_final_d.pt")
|
||
torch.save(ema.model.state_dict(), save_dir / f"{run_name}_final_ema.pt")
|
||
history["train_time_s"] = time.time() - t_start
|
||
return history
|
||
|
||
|
||
# ────────────────────────────────────────────────────────────────────────────
|
||
# Phase 3 — VAE (3.1 MSE+KL · 3.2 +Perceptual · 3.3 +PatchGAN)
|
||
# ────────────────────────────────────────────────────────────────────────────
|
||
|
||
def _save_vae_samples(
|
||
vae,
|
||
samples_dir: Path,
|
||
epoch: int,
|
||
*,
|
||
fixed_z: torch.Tensor,
|
||
fixed_real: torch.Tensor,
|
||
device,
|
||
) -> None:
|
||
"""Save prior samples and a real-vs-reconstruction grid side by side."""
|
||
samples_dir.mkdir(parents=True, exist_ok=True)
|
||
vae.eval()
|
||
with torch.no_grad():
|
||
prior = vae.decode(fixed_z.to(device))
|
||
prior = (prior.clamp(-1, 1) + 1.0) / 2.0
|
||
save_image(prior, samples_dir / f"epoch_{epoch:04d}.png", nrow=4)
|
||
|
||
recon, _, _ = vae(fixed_real.to(device))
|
||
recon = (recon.clamp(-1, 1) + 1.0) / 2.0
|
||
real = (fixed_real.to(device) + 1.0) / 2.0
|
||
# Interleave real / reconstruction pairs
|
||
pairs = torch.stack([real, recon], dim=1).flatten(0, 1)
|
||
save_image(pairs, samples_dir / f"epoch_{epoch:04d}_recon.png", nrow=4)
|
||
vae.train()
|
||
|
||
|
||
def train_vae(
|
||
vae,
|
||
train_dataset,
|
||
cfg: dict,
|
||
*,
|
||
save_dir,
|
||
run_name: str,
|
||
device: str = "cuda",
|
||
) -> dict:
|
||
"""VAE training loop covering Phase 3.1 – 3.3 and Phase 5.
|
||
|
||
Config toggles:
|
||
lambda_perceptual > 0 → VGG-16 perceptual loss (Phase 3.2+)
|
||
lambda_adversarial > 0 → PatchGAN hinge adversarial loss (Phase 3.3)
|
||
|
||
Loss: L = L_mse + λ_perc·L_vgg + λ_adv·L_adv + β_kl·L_kl
|
||
|
||
KL is computed as mean over latent dimensions (scale-invariant), so
|
||
beta_kl is comparable across different latent_dim values.
|
||
|
||
AMP is intentionally disabled for VAE training — mixed-precision float16
|
||
overflows when the KL divergence spikes, producing NaN cascades that
|
||
corrupt the model irrecoverably. All VAE + perceptual + PatchGAN
|
||
computation runs in float32.
|
||
"""
|
||
device = torch.device(device if torch.cuda.is_available() else "cpu")
|
||
vae = vae.to(device)
|
||
|
||
n_vae = sum(p.numel() for p in vae.parameters() if p.requires_grad)
|
||
print(f"VAE: {n_vae:,} params")
|
||
|
||
epochs = cfg["epochs"]
|
||
batch_size = cfg["batch_size"]
|
||
lr = cfg.get("lr", 1e-3)
|
||
latent_dim = cfg.get("latent_dim", 256)
|
||
beta_kl = cfg.get("beta_kl", 1.0)
|
||
lambda_perceptual = cfg.get("lambda_perceptual", 0.0)
|
||
lambda_adversarial = cfg.get("lambda_adversarial", 0.0)
|
||
lr_d = cfg.get("lr_d", 1e-4)
|
||
grad_clip = cfg.get("grad_clip", 1.0)
|
||
ema_decay = cfg.get("ema_decay", 0.9999)
|
||
sample_interval = cfg.get("sample_interval", 10)
|
||
fid_interval = cfg.get("fid_interval", 25)
|
||
fid_n_real = cfg.get("fid_n_real", 5000)
|
||
|
||
use_perceptual = lambda_perceptual > 0
|
||
use_adversarial = lambda_adversarial > 0
|
||
|
||
loader = DataLoader(
|
||
train_dataset, batch_size=batch_size, shuffle=True,
|
||
num_workers=cfg.get("num_workers", min(4, os.cpu_count() or 1)),
|
||
pin_memory=(device.type == "cuda"), drop_last=True,
|
||
)
|
||
|
||
opt_vae = torch.optim.Adam(vae.parameters(), lr=lr)
|
||
# AMP disabled — float16 overflows on KL spikes, causing NaN cascades
|
||
use_amp = False
|
||
|
||
# KL warmup: linearly ramp beta_kl from 0 to target over first 20% of training
|
||
kl_warmup_epochs = max(1, epochs // 5)
|
||
|
||
# Linear LR decay from epoch epochs//2 to epochs
|
||
decay_start = epochs // 2
|
||
sched_vae = torch.optim.lr_scheduler.LambdaLR(
|
||
opt_vae, lr_lambda=lambda ep: max(0.0, 1.0 - max(ep - decay_start, 0) / max(epochs - decay_start, 1)))
|
||
sched_d = None # set below if adversarial
|
||
|
||
# ── Optional components ───────────────────────────────────────────────
|
||
perc_fn = None
|
||
patchgan = None
|
||
opt_d = None
|
||
|
||
if use_perceptual:
|
||
from src.training.perceptual import PerceptualLoss
|
||
perc_fn = PerceptualLoss().to(device).float()
|
||
print("Perceptual loss: VGG-16 relu1_2 + relu2_2 + relu3_3")
|
||
|
||
if use_adversarial:
|
||
from src.models.patchgan import PatchGANDiscriminator, hinge_d_loss, hinge_g_loss
|
||
patchgan = PatchGANDiscriminator(
|
||
ndf=cfg.get("ndf_patch", 64),
|
||
image_size=cfg.get("image_size", 64),
|
||
).to(device).float()
|
||
opt_d = torch.optim.Adam(patchgan.parameters(), lr=lr_d, betas=(0.5, 0.999))
|
||
sched_d = torch.optim.lr_scheduler.LambdaLR(
|
||
opt_d, lr_lambda=lambda ep: max(0.0, 1.0 - max(ep - decay_start, 0) / max(epochs - decay_start, 1)))
|
||
n_d = sum(p.numel() for p in patchgan.parameters())
|
||
print(f"PatchGAN: {n_d:,} params")
|
||
else:
|
||
hinge_d_loss = hinge_g_loss = None # never called
|
||
|
||
# ── Fixed seeds for consistent visualisation ──────────────────────────
|
||
fixed_z = torch.randn(16, latent_dim, device=device)
|
||
# Grab first 16 real images from the loader for reconstruction tracking
|
||
_it = iter(loader)
|
||
fixed_real = next(_it)[:16].cpu()
|
||
|
||
ema = EMA(vae, decay=ema_decay)
|
||
|
||
save_dir = Path(save_dir)
|
||
save_dir.mkdir(parents=True, exist_ok=True)
|
||
samples_dir = save_dir.parent / "samples" / run_name
|
||
|
||
fid_eval = FIDEvaluator(train_dataset, n_real=fid_n_real, device=str(device),
|
||
num_workers=cfg.get("num_workers", min(4, os.cpu_count() or 1)))
|
||
|
||
history = {
|
||
"recon_loss": [], "kl_loss": [], "perc_loss": [],
|
||
"adv_g_loss": [], "adv_d_loss": [], "fid": {},
|
||
}
|
||
best_fid = float("inf")
|
||
nan_skipped = 0
|
||
print(
|
||
f"Device: {device} AMP: disabled (float32) Batches/epoch: {len(loader)}"
|
||
f" β_kl={beta_kl} (warmup {kl_warmup_epochs}ep) λ_perc={lambda_perceptual}"
|
||
f" λ_adv={lambda_adversarial}"
|
||
)
|
||
|
||
t_start = time.time()
|
||
|
||
for epoch in range(1, epochs + 1):
|
||
vae.train()
|
||
if patchgan is not None:
|
||
patchgan.train()
|
||
|
||
recon_sum = kl_sum = perc_sum = adv_g_sum = adv_d_sum = 0.0
|
||
n_batches = 0
|
||
|
||
for real in tqdm(loader, desc=f"Epoch {epoch}/{epochs}", leave=False):
|
||
real = real.to(device).float()
|
||
|
||
# KL warmup: ramp from 0 to beta_kl over kl_warmup_epochs
|
||
current_beta = beta_kl * min(1.0, epoch / kl_warmup_epochs)
|
||
|
||
# ── VAE forward (float32, no AMP) ────────────────────────────
|
||
recon, mu, log_var = vae(real)
|
||
mse = F.mse_loss(recon, real)
|
||
|
||
# KL divergence: mean over latent dims (scale-invariant w.r.t. latent_dim)
|
||
kl = (-0.5 * (1 + log_var - mu.pow(2) - log_var.exp())).mean()
|
||
|
||
perc = perc_fn(recon, real) if use_perceptual else real.new_zeros(1).squeeze()
|
||
vae_loss = mse + current_beta * kl + lambda_perceptual * perc
|
||
|
||
# ── NaN/Inf guard ────────────────────────────────────────────
|
||
if not torch.isfinite(vae_loss):
|
||
nan_skipped += 1
|
||
opt_vae.zero_grad()
|
||
continue
|
||
|
||
# ── PatchGAN discriminator step ───────────────────────────────
|
||
adv_d = real.new_zeros(1).squeeze()
|
||
if use_adversarial:
|
||
opt_d.zero_grad()
|
||
d_real = patchgan(real)
|
||
d_fake = patchgan(recon.detach())
|
||
adv_d = hinge_d_loss(d_real, d_fake)
|
||
if torch.isfinite(adv_d):
|
||
adv_d.backward()
|
||
torch.nn.utils.clip_grad_norm_(patchgan.parameters(), grad_clip)
|
||
opt_d.step()
|
||
|
||
# ── PatchGAN generator adversarial loss ───────────────────────
|
||
adv_g = real.new_zeros(1).squeeze()
|
||
if use_adversarial:
|
||
adv_g = hinge_g_loss(patchgan(recon))
|
||
vae_loss = vae_loss + lambda_adversarial * adv_g
|
||
|
||
# ── VAE backward ──────────────────────────────────────────────
|
||
opt_vae.zero_grad()
|
||
vae_loss.backward()
|
||
torch.nn.utils.clip_grad_norm_(vae.parameters(), grad_clip)
|
||
opt_vae.step()
|
||
ema.update(vae)
|
||
|
||
recon_sum += mse.item()
|
||
kl_sum += kl.item()
|
||
perc_sum += perc.item()
|
||
adv_g_sum += adv_g.item()
|
||
adv_d_sum += adv_d.item()
|
||
n_batches += 1
|
||
|
||
avg_r = recon_sum / max(n_batches, 1)
|
||
avg_k = kl_sum / max(n_batches, 1)
|
||
avg_p = perc_sum / max(n_batches, 1)
|
||
avg_g = adv_g_sum / max(n_batches, 1)
|
||
avg_d = adv_d_sum / max(n_batches, 1)
|
||
history["recon_loss"].append(avg_r)
|
||
history["kl_loss"].append(avg_k)
|
||
history["perc_loss"].append(avg_p)
|
||
history["adv_g_loss"].append(avg_g)
|
||
history["adv_d_loss"].append(avg_d)
|
||
|
||
print(
|
||
f"[{epoch:03d}/{epochs}] "
|
||
f"MSE: {avg_r:.4f} KL: {avg_k:.2f} β={current_beta:.6f} "
|
||
f"Perc: {avg_p:.4f} AdvG: {avg_g:.4f} AdvD: {avg_d:.4f}"
|
||
f" (NaN skipped: {nan_skipped})"
|
||
)
|
||
|
||
if epoch % sample_interval == 0:
|
||
_save_vae_samples(
|
||
ema.model, samples_dir, epoch,
|
||
fixed_z=fixed_z, fixed_real=fixed_real, device=device,
|
||
)
|
||
|
||
if epoch % fid_interval == 0:
|
||
ema.model.eval()
|
||
with torch.no_grad():
|
||
fake_imgs = torch.cat([
|
||
ema.model.sample(64, device)
|
||
for _ in range(fid_n_real // 64 + 1)
|
||
])[:fid_n_real]
|
||
fid_score = fid_eval.compute(fake_imgs)
|
||
history["fid"][epoch] = fid_score
|
||
print(f" FID @ epoch {epoch}: {fid_score:.2f}")
|
||
|
||
if fid_score < best_fid:
|
||
best_fid = fid_score
|
||
torch.save(vae.state_dict(), save_dir / f"{run_name}_best_vae.pt")
|
||
torch.save(ema.model.state_dict(), save_dir / f"{run_name}_best_ema.pt")
|
||
|
||
sched_vae.step()
|
||
if sched_d is not None:
|
||
sched_d.step()
|
||
|
||
torch.save(vae.state_dict(), save_dir / f"{run_name}_final_vae.pt")
|
||
torch.save(ema.model.state_dict(), save_dir / f"{run_name}_final_ema.pt")
|
||
if patchgan is not None:
|
||
torch.save(patchgan.state_dict(), save_dir / f"{run_name}_final_patchgan.pt")
|
||
history["train_time_s"] = time.time() - t_start
|
||
print(f"Total NaN-skipped batches: {nan_skipped}")
|
||
return history
|
||
|
||
|
||
# ────────────────────────────────────────────────────────────────────────────
|
||
# Phase 4 — DDPM (4.1 linear·ε · 4.2 cosine·ε · 4.3 cosine·v · 4.4 wider)
|
||
# ────────────────────────────────────────────────────────────────────────────
|
||
|
||
def train_ddpm(
|
||
model,
|
||
train_dataset,
|
||
cfg: dict,
|
||
*,
|
||
save_dir,
|
||
run_name: str,
|
||
device: str = "cuda",
|
||
) -> dict:
|
||
"""DDPM training loop (Ho et al., 2020) covering Phase 4.1 – 4.4.
|
||
|
||
Config keys:
|
||
noise_schedule — "linear" (4.1) or "cosine" (4.2+)
|
||
pred_type — "eps" (4.1–4.2) or "v" (4.3+)
|
||
T — diffusion timesteps (default 1000)
|
||
base_ch / ch_mult / attn_resolutions — U-Net capacity (see unet.py)
|
||
ddim_steps — DDIM steps for FID evaluation (default 100)
|
||
"""
|
||
from src.training.diffusion import (
|
||
linear_betas, cosine_betas, make_alpha_bars,
|
||
diffusion_loss, ddim_sample,
|
||
)
|
||
|
||
device = torch.device(device if torch.cuda.is_available() else "cpu")
|
||
model = model.to(device)
|
||
|
||
n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||
print(f"U-Net: {n_params:,} params")
|
||
|
||
epochs = cfg["epochs"]
|
||
batch_size = cfg["batch_size"]
|
||
lr = cfg.get("lr", 2e-4)
|
||
T = cfg.get("T", 1000)
|
||
noise_schedule = cfg.get("noise_schedule", "linear")
|
||
pred_type = cfg.get("pred_type", "eps")
|
||
ddim_steps = cfg.get("ddim_steps", 100)
|
||
image_size = cfg.get("image_size", 64)
|
||
ema_decay = cfg.get("ema_decay", 0.9999)
|
||
sample_interval = cfg.get("sample_interval", 10)
|
||
fid_interval = cfg.get("fid_interval", 25)
|
||
fid_n_real = cfg.get("fid_n_real", 5000)
|
||
|
||
# Build noise schedule and register on device
|
||
betas = (cosine_betas(T) if noise_schedule == "cosine" else linear_betas(T)).to(device)
|
||
alpha_bars = make_alpha_bars(betas) # on device
|
||
|
||
loader = DataLoader(
|
||
train_dataset, batch_size=batch_size, shuffle=True,
|
||
num_workers=cfg.get("num_workers", min(4, os.cpu_count() or 1)),
|
||
pin_memory=(device.type == "cuda"), drop_last=True,
|
||
)
|
||
opt = torch.optim.AdamW(model.parameters(), lr=lr)
|
||
|
||
use_amp = device.type == "cuda"
|
||
scaler = _GradScaler("cuda", enabled=use_amp)
|
||
|
||
ema = EMA(model, decay=ema_decay)
|
||
|
||
# Fixed noise for sample visualisation (same latents across epochs)
|
||
fixed_noise = torch.randn(16, 3, image_size, image_size, device=device)
|
||
|
||
save_dir = Path(save_dir)
|
||
save_dir.mkdir(parents=True, exist_ok=True)
|
||
samples_dir = save_dir.parent / "samples" / run_name
|
||
|
||
fid_eval = FIDEvaluator(train_dataset, n_real=fid_n_real, device=str(device),
|
||
num_workers=cfg.get("num_workers", min(4, os.cpu_count() or 1)))
|
||
|
||
history = {"loss": [], "fid": {}}
|
||
best_fid = float("inf")
|
||
print(
|
||
f"Device: {device} AMP: {use_amp} Batches/epoch: {len(loader)}"
|
||
f" T={T} schedule={noise_schedule} pred={pred_type} ddim_steps={ddim_steps}"
|
||
)
|
||
|
||
# Linear LR decay from epoch epochs//2 to epochs
|
||
decay_start = epochs // 2
|
||
sched = torch.optim.lr_scheduler.LambdaLR(
|
||
opt, lr_lambda=lambda ep: max(0.0, 1.0 - max(ep - decay_start, 0) / max(epochs - decay_start, 1)))
|
||
|
||
t_start = time.time()
|
||
|
||
for epoch in range(1, epochs + 1):
|
||
model.train()
|
||
loss_sum = 0.0
|
||
n_batches = 0
|
||
|
||
for x0 in tqdm(loader, desc=f"Epoch {epoch}/{epochs}", leave=False):
|
||
x0 = x0.to(device)
|
||
t = torch.randint(0, T, (x0.size(0),), device=device)
|
||
|
||
with _autocast("cuda", enabled=use_amp):
|
||
loss = diffusion_loss(model, x0, t, alpha_bars, pred_type)
|
||
|
||
opt.zero_grad()
|
||
scaler.scale(loss).backward()
|
||
scaler.unscale_(opt)
|
||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||
scaler.step(opt)
|
||
scaler.update()
|
||
ema.update(model)
|
||
|
||
loss_sum += loss.item()
|
||
n_batches += 1
|
||
|
||
avg_loss = loss_sum / n_batches
|
||
history["loss"].append(avg_loss)
|
||
print(f"[{epoch:03d}/{epochs}] Loss: {avg_loss:.5f}")
|
||
|
||
if epoch % sample_interval == 0:
|
||
samples_dir.mkdir(parents=True, exist_ok=True)
|
||
ema.model.eval()
|
||
with torch.no_grad():
|
||
# Quick visualisation: denoise fixed_noise via DDIM
|
||
imgs = ddim_sample(
|
||
ema.model, 16, image_size, alpha_bars,
|
||
n_steps=50, pred_type=pred_type, device=str(device), batch_size=16,
|
||
)
|
||
imgs = (imgs.clamp(-1, 1) + 1.0) / 2.0
|
||
save_image(imgs, samples_dir / f"epoch_{epoch:04d}.png", nrow=4)
|
||
|
||
if epoch % fid_interval == 0:
|
||
ema.model.eval()
|
||
fake_imgs = ddim_sample(
|
||
ema.model, fid_n_real, image_size, alpha_bars,
|
||
n_steps=ddim_steps, pred_type=pred_type,
|
||
device=str(device), batch_size=32,
|
||
)
|
||
fid_score = fid_eval.compute(fake_imgs)
|
||
history["fid"][epoch] = fid_score
|
||
print(f" FID @ epoch {epoch}: {fid_score:.2f}")
|
||
|
||
if fid_score < best_fid:
|
||
best_fid = fid_score
|
||
torch.save(model.state_dict(), save_dir / f"{run_name}_best_unet.pt")
|
||
torch.save(ema.model.state_dict(), save_dir / f"{run_name}_best_ema.pt")
|
||
|
||
sched.step()
|
||
|
||
torch.save(model.state_dict(), save_dir / f"{run_name}_final_unet.pt")
|
||
torch.save(ema.model.state_dict(), save_dir / f"{run_name}_final_ema.pt")
|
||
history["train_time_s"] = time.time() - t_start
|
||
return history
|