diff options
Diffstat (limited to 'src/train_diffusion.py')
| -rw-r--r-- | src/train_diffusion.py | 130 |
1 files changed, 103 insertions, 27 deletions
diff --git a/src/train_diffusion.py b/src/train_diffusion.py index 2ec4d60..d0777bc 100644 --- a/src/train_diffusion.py +++ b/src/train_diffusion.py @@ -3,16 +3,13 @@ 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.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_dataloader, get_mnist_dataset +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, CHECKPOINT_DIR, PLOTS_DIR + DIFFU_STEPS, NB_CHANNEL, IMG_SIZE, NB_LABEL ) import os import torch @@ -28,11 +25,14 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) def train_diffusion( - epochs: int = EPOCHS, + 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, ): @@ -45,9 +45,17 @@ def train_diffusion( 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 @@ -69,20 +77,36 @@ def train_diffusion( optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) criterion = nn.MSELoss() - # Get dataloader - dataloader = get_diffusion_dataloader( + # 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 - for epoch in range(1, epochs + 1): + epoch = 0 + while epochs is None or epoch < epochs: + epoch += 1 + model.train() epoch_loss = 0.0 - for batch_idx, (xt, t, vec, target) in enumerate(track(dataloader, description=f"Epoch {epoch}/{epochs}")): + 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 @@ -99,25 +123,68 @@ def train_diffusion( # 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) + mlflow.log_metric("train/loss_step", avg_loss, step=global_step) - avg_loss = epoch_loss / len(dataloader) - mlflow.log_metric("epoch_loss", avg_loss, step=epoch) + avg_loss = epoch_loss / len(train_loader) + mlflow.log_metric("train/loss", avg_loss, step=epoch) - # Save checkpoint every 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") - # Generate and log sample images + # 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, predict_x0) + evaluate_and_log(model, epoch, test_loader, 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.""" +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(): @@ -142,10 +209,13 @@ def evaluate_and_log(model: nn.Module, epoch: int, predict_x0: bool = True): fakes = np.concatenate(fakes) - # Get real samples for comparison - dataset = get_mnist_dataset(train=True) + # Get real samples from test split n_samples = len(fakes) - reals = torch.stack([dataset[i][0] for i in range(n_samples)]).numpy() + 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 @@ -153,9 +223,9 @@ def evaluate_and_log(model: nn.Module, epoch: int, predict_x0: bool = True): 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) + 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) @@ -176,17 +246,20 @@ def evaluate_and_log(model: nn.Module, epoch: int, predict_x0: bool = True): fig.update_xaxes(showticklabels=False, showgrid=False, zeroline=False) fig.update_yaxes(showticklabels=False, showgrid=False, zeroline=False) - mlflow.log_figure(fig, f"samples/epoch_{epoch:03d}.png") + 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=EPOCHS, help="Number of epochs") + 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") @@ -199,6 +272,9 @@ if __name__ == "__main__": 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, ) |