""" Training Script for Conditional Diffusion Model Trains diffusion model conditioned on cosmological parameters (Omega_m, sigma_8) Changes from original: - EMA weights are now applied before validation and sampling - Training args are saved to args.txt for evaluation script - Fixed --normalize_labels and --use_ddim flags (were un-disableable) - Added mixed-precision (AMP) training support - Fixed loss averaging to be per-sample rather than per-batch - Added weights_only=True to torch.load for security """ import torch import torch.optim as optim import numpy as np import os import argparse import json import random from tqdm import tqdm import time from unet_conditional import ConditionalUNet from diffusion_conditional import GaussianDiffusion, ConditionalDiffusionModel from dataset_conditional import get_conditional_dataloaders import matplotlib.pyplot as plt # Weights & Biases (optional) try: import wandb WANDB_AVAILABLE = True except ImportError: WANDB_AVAILABLE = False print("Warning: wandb not available. Install with: pip install wandb") class EMA: """Exponential Moving Average for model parameters""" def __init__(self, model, decay=0.9999): self.model = model self.decay = decay self.shadow = {} for name, param in model.named_parameters(): if param.requires_grad: self.shadow[name] = param.data.clone() def update(self): for name, param in self.model.named_parameters(): if param.requires_grad: self.shadow[name] = self.decay * self.shadow[name] + (1 - self.decay) * param.data def apply_shadow(self): self.backup = {name: param.data.clone() for name, param in self.model.named_parameters() if param.requires_grad} for name, param in self.model.named_parameters(): if param.requires_grad: param.data = self.shadow[name] def restore(self): for name, param in self.model.named_parameters(): if param.requires_grad: param.data = self.backup[name] self.backup = {} def train_epoch(model, dataloader, optimizer, device, epoch, ema=None, use_wandb=False, scaler=None): model.train() total_loss = 0.0 total_samples = 0 pbar = tqdm(dataloader, desc=f'Epoch {epoch}') for batch_idx, (images, labels) in enumerate(pbar): images = images.to(device) labels = labels.to(device) batch_size = images.shape[0] optimizer.zero_grad() if scaler is not None: with torch.amp.autocast('cuda'): loss = model.get_loss(images, labels) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update() else: loss = model.get_loss(images, labels) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() if ema is not None: ema.update() total_loss += loss.item() * batch_size total_samples += batch_size pbar.set_postfix({'loss': f'{loss.item():.4f}'}) if use_wandb and batch_idx % 10 == 0: wandb.log({'batch_loss': loss.item(), 'epoch': epoch, 'batch': epoch * len(dataloader) + batch_idx}) return total_loss / total_samples def validate(model, dataloader, device): model.eval() total_loss = 0.0 total_samples = 0 with torch.no_grad(): for images, labels in tqdm(dataloader, desc='Validating'): images = images.to(device) labels = labels.to(device) batch_size = images.shape[0] loss = model.get_loss(images, labels) total_loss += loss.item() * batch_size total_samples += batch_size return total_loss / total_samples def save_checkpoint(model, optimizer, ema, epoch, loss, save_dir, is_best=False, last_improvement_epoch=None, scheduler=None): checkpoint = { 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': loss, } if ema is not None: checkpoint['ema_shadow'] = ema.shadow if last_improvement_epoch is not None: checkpoint['last_improvement_epoch'] = last_improvement_epoch if scheduler is not None: checkpoint['scheduler_state_dict'] = scheduler.state_dict() torch.save(checkpoint, os.path.join(save_dir, 'checkpoint_latest.pt')) if is_best: torch.save(checkpoint, os.path.join(save_dir, 'best_model.pt')) print(f"Saved best model at epoch {epoch+1}") if (epoch + 1) % 20 == 0: torch.save(checkpoint, os.path.join(save_dir, f'checkpoint_epoch_{epoch+1}.pt')) print(f"Saved checkpoint at epoch {epoch+1}") def sample_images(model, diffusion, device, save_path, test_labels, ema=None, n_samples=8, epoch=0, use_ddim=True, ddim_steps=50, use_wandb=False): # Apply EMA weights for sampling if ema is not None: ema.apply_shadow() model.eval() labels = test_labels[:n_samples].to(device) with torch.no_grad(): samples = diffusion.sample( model, labels=labels, channels=1, height=256, width=256, device=device, progress=True, use_ddim=use_ddim, ddim_steps=ddim_steps, eta=0.0 ) # Restore original weights after sampling if ema is not None: ema.restore() n_cols = min(n_samples, 4) n_rows = (n_samples + n_cols - 1) // n_cols fig, axes = plt.subplots(n_rows, n_cols, figsize=(4.5 * n_cols, 4.5 * n_rows)) if n_rows == 1 and n_cols == 1: axes = np.array([[axes]]) elif n_rows == 1: axes = axes[np.newaxis, :] elif n_cols == 1: axes = axes[:, np.newaxis] for i in range(n_rows * n_cols): ax = axes[i // n_cols, i % n_cols] if i < n_samples: img = samples[i, 0].cpu().numpy() label_vals = labels[i].cpu().tolist() label_str = ", ".join(f"{v:.2f}" for v in label_vals) ax.imshow(img, cmap='gray', vmin=-1, vmax=1) ax.set_title(label_str, fontsize=10) ax.axis('off') plt.suptitle(f'Generated Samples - Epoch {epoch}', fontsize=14) plt.tight_layout() plt.savefig(save_path, dpi=150, bbox_inches='tight') if use_wandb: wandb.log({'generated_samples': wandb.Image(save_path), 'epoch': epoch}) plt.close() print(f"Saved samples to {save_path}") def save_training_args(args, output_dir): """Save training arguments so the evaluation script can reconstruct the model.""" args_path = os.path.join(output_dir, 'args.txt') with open(args_path, 'w') as f: for key, value in vars(args).items(): f.write(f"{key}: {value}\n") # Also save as JSON for robust parsing args_json_path = os.path.join(output_dir, 'args.json') with open(args_json_path, 'w') as f: json.dump(vars(args), f, indent=2) print(f"Saved training args to {args_path} and {args_json_path}") def main(): parser = argparse.ArgumentParser(description='Train Conditional Diffusion Model') # Model parser.add_argument('--label_dim', type=int, default=2) parser.add_argument('--base_channels', type=int, default=64) parser.add_argument('--channel_multipliers', type=int, nargs='+', default=[1, 2, 4, 8]) parser.add_argument('--attention_levels', type=int, nargs='+', default=[2, 3]) parser.add_argument('--dropout', type=float, default=0.1) # Diffusion parser.add_argument('--timesteps', type=int, default=1500) parser.add_argument('--beta_start', type=float, default=1e-4) parser.add_argument('--beta_end', type=float, default=0.02) parser.add_argument('--schedule_type', type=str, default='linear') # Training parser.add_argument('--epochs', type=int, default=100) parser.add_argument('--batch_size', type=int, default=8) parser.add_argument('--lr', type=float, default=2e-4) parser.add_argument('--ema_decay', type=float, default=0.9999) parser.add_argument('--num_workers', type=int, default=4) parser.add_argument('--early_stop_patience', type=int, default=30) parser.add_argument('--use_amp', action='store_true', default=False, help='Enable mixed-precision training (recommended for GPU)') # Data parser.add_argument('--data_dir', type=str, default='./data/params_2', help='Data directory (relative to repo root)') # FIX: Use BooleanOptionalAction so --no-normalize-labels works parser.add_argument('--normalize_labels', action=argparse.BooleanOptionalAction, default=True) # Output parser.add_argument('--output_dir', type=str, default='outputs_conditional') parser.add_argument('--resume', type=str, default='') parser.add_argument( '--resume_refresh_scheduler', action='store_true', help='On resume, rebuild cosine LR scheduler for --epochs (last_epoch=start-1) instead of loading saved scheduler; use when extending training beyond the original epoch count', ) parser.add_argument('--sample_every', type=int, default=10) # FIX: Use BooleanOptionalAction so --no-use-ddim works parser.add_argument('--use_ddim', action=argparse.BooleanOptionalAction, default=True) parser.add_argument('--ddim_steps', type=int, default=50) # WandB parser.add_argument('--use_wandb', action='store_true', default=False) parser.add_argument('--wandb_project', type=str, default='ddpm_cosmology') parser.add_argument('--wandb_entity', type=str, default='') parser.add_argument('--wandb_run_name', type=str, default='') args = parser.parse_args() # Reproducibility seed = 42 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False # WandB use_wandb = args.use_wandb and WANDB_AVAILABLE if use_wandb: run_name = args.wandb_run_name or f"conditional_diffusion_{time.strftime('%Y%m%d_%H%M%S')}" wandb.init(project=args.wandb_project, entity=args.wandb_entity or None, name=run_name, config=vars(args)) print(f"W&B run: {run_name}") # Directories timestamp = time.strftime("%Y%m%d_%H%M%S") output_dir = f"{args.output_dir}_{timestamp}" os.makedirs(output_dir, exist_ok=True) os.makedirs(os.path.join(output_dir, 'checkpoints'), exist_ok=True) os.makedirs(os.path.join(output_dir, 'samples'), exist_ok=True) # Save training args for evaluation save_training_args(args, output_dir) # Device device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Using device: {device}") # AMP scaler scaler = torch.amp.GradScaler('cuda') if args.use_amp and torch.cuda.is_available() else None if scaler: print("Mixed-precision training enabled (AMP)") # Data print("\nLoading data...") train_loader, val_loader, test_loader = get_conditional_dataloaders( data_dir=args.data_dir, batch_size=args.batch_size, num_workers=args.num_workers, normalize_labels=args.normalize_labels ) _, test_labels = next(iter(test_loader)) # Model print("\nCreating model...") unet = ConditionalUNet( in_channels=1, out_channels=1, label_dim=args.label_dim, base_channels=args.base_channels, channel_multipliers=args.channel_multipliers, attention_levels=args.attention_levels, dropout=args.dropout ) diffusion = GaussianDiffusion( timesteps=args.timesteps, beta_start=args.beta_start, beta_end=args.beta_end, schedule_type=args.schedule_type ) model = ConditionalDiffusionModel(unet, diffusion).to(device) print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}") optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01) ema = EMA(model, decay=args.ema_decay) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) # Resume start_epoch = 0 best_val_loss = float('inf') last_improvement_epoch = -1 if args.resume: print(f"Resuming from {args.resume}") checkpoint = torch.load(args.resume, map_location=device, weights_only=False) model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) if 'ema_shadow' in checkpoint: ema.shadow = checkpoint['ema_shadow'] start_epoch = checkpoint['epoch'] + 1 best_val_loss = checkpoint.get('loss', float('inf')) last_improvement_epoch = checkpoint.get('last_improvement_epoch', -1) if args.resume_refresh_scheduler: scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=args.epochs, last_epoch=start_epoch - 1 ) print( f"Rebuilt LR scheduler for extended run: T_max={args.epochs}, " f"resume at epoch {start_epoch + 1} (last_epoch={start_epoch - 1})" ) elif 'scheduler_state_dict' in checkpoint: scheduler.load_state_dict(checkpoint['scheduler_state_dict']) # Training print("\nStarting training...") losses = {'train': [], 'val': []} for epoch in range(start_epoch, args.epochs): train_loss = train_epoch(model, train_loader, optimizer, device, epoch, ema, use_wandb, scaler=scaler) # Apply EMA weights for validation if ema is not None: ema.apply_shadow() val_loss = validate(model, val_loader, device) if ema is not None: ema.restore() losses['train'].append(train_loss) losses['val'].append(val_loss) scheduler.step() if use_wandb: wandb.log({ 'epoch': epoch + 1, 'train_loss': train_loss, 'val_loss': val_loss, 'learning_rate': optimizer.param_groups[0]['lr'] }) print(f"\nEpoch {epoch+1}/{args.epochs} | Train: {train_loss:.6f} | Val: {val_loss:.6f} | LR: {optimizer.param_groups[0]['lr']:.6e}") is_best = val_loss < best_val_loss if is_best: best_val_loss = val_loss last_improvement_epoch = epoch save_checkpoint(model, optimizer, ema, epoch, val_loss, os.path.join(output_dir, 'checkpoints'), is_best=is_best, last_improvement_epoch=last_improvement_epoch, scheduler=scheduler) # Early stopping if epoch - last_improvement_epoch >= args.early_stop_patience: print(f"Early stopping at epoch {epoch+1}") break # Samples (with EMA weights) if (epoch + 1) % args.sample_every == 0: sample_path = os.path.join(output_dir, 'samples', f'samples_epoch_{epoch+1}.png') sample_images(model, diffusion, device, sample_path, test_labels, ema=ema, epoch=epoch+1, use_ddim=args.use_ddim, ddim_steps=args.ddim_steps, use_wandb=use_wandb) # Loss plot if (epoch + 1) % 5 == 0: plt.figure(figsize=(10, 5)) plt.plot(losses['train'], label='Train Loss') plt.plot(losses['val'], label='Val Loss') plt.yscale('log') plt.xlabel('Epoch') plt.ylabel('Loss') plt.title('Training Progress') plt.legend() plt.grid(True, alpha=0.3) plt.savefig(os.path.join(output_dir, 'losses.png'), dpi=150) plt.close() print(f"\nTraining completed! Best val loss: {best_val_loss:.6f}") print(f"Results saved to: {output_dir}") if use_wandb: wandb.finish() if __name__ == '__main__': main()