#!/usr/bin/env python3 """ WP3 -- train the grid-based model. python scripts/03_train.py --config Configs/train_grid_12keV.yaml python scripts/03_train.py --config ... --epochs 3 --tag smoke Every epoch the validation pass reports both the loss and the *detection* metrics (efficiency / purity / position RMS), because the loss alone says nothing about the working point. """ import argparse import json import os import sys import time from datetime import datetime from pathlib import Path import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import torch from omegaconf import OmegaConf from torch.utils.data import DataLoader HERE = Path(__file__).resolve().parent sys.path.insert(0, str(HERE.parent / "src")) from datasets import GridFrameDataset, collate_with_photons # noqa: E402 from decode import decode # noqa: E402 from losses import GridLoss # noqa: E402 from metrics import DetectionStats # noqa: E402 from models import build_model, valid_mask # noqa: E402 LOSS_PARTS = ("objectness", "offset", "energy") def shard_files(folder, split, shard_range): first, last = shard_range return [os.path.join(folder, split, f"shard_{index:02d}.h5") for index in range(int(first), int(last) + 1)] def prepare_output_dir(conf, tag=None): date = datetime.now().strftime("%y%m%d") base = Path(conf.experiment.results_dir) run_index = 0 while True: name = (f"{date}_{conf.experiment.name}" + (f"_{tag}" if tag else "") + f"_{run_index:02d}") output_dir = base / name if not output_dir.exists(): break run_index += 1 (output_dir / "Models").mkdir(parents=True) (output_dir / "Plots").mkdir(parents=True) OmegaConf.save(conf, output_dir / "config.yaml") return output_dir def run_epoch(model, loader, loss_function, mask, optimizer=None, grad_scaler=None, use_amp=False): is_training = optimizer is not None model.train(is_training) totals = dict(loss=0.0, objectness=0.0, offset=0.0, energy=0.0) n_samples = 0 for batch in loader: frames, objectness, offset, energy = [t.cuda(non_blocking=True) for t in batch[:4]] with torch.autocast("cuda", enabled=use_amp): prediction = model(frames) loss, parts = loss_function(prediction.float(), objectness, offset, energy, mask) if is_training: optimizer.zero_grad(set_to_none=True) if grad_scaler is not None: grad_scaler.scale(loss).backward() grad_scaler.step(optimizer) grad_scaler.update() else: loss.backward() optimizer.step() batch_size = frames.shape[0] totals["loss"] += loss.item() * batch_size for part in LOSS_PARTS: totals[part] += parts[part].item() * batch_size n_samples += batch_size return {key: value / max(n_samples, 1) for key, value in totals.items()} @torch.no_grad() def run_detection_metrics(model, loader, conf, halo, frame_size, use_amp=False): model.eval() stats = DetectionStats(pixel_size_um=float(conf.eval.pixel_size_um), match_radius=float(conf.eval.match_radius), halo=halo, frame_size=frame_size) for frames, objectness, offset, energy, photons in loader: frames = frames.cuda(non_blocking=True) with torch.autocast("cuda", enabled=use_amp): prediction = model(frames) points = decode(prediction.float(), threshold=float(conf.eval.threshold), merge_distance=float(conf.eval.merge_distance), halo=halo, predict_energy=bool(conf.model.predict_energy), energy_scale=float(conf.data.energy_scale)) stats.update([frame_points.numpy() for frame_points in points], [frame_truth.numpy() for frame_truth in photons]) return stats.summary() def plot_curves(history, output_dir): epochs = np.arange(1, len(history["train_loss"]) + 1) figure, axes = plt.subplots(1, 3, figsize=(15, 4.2)) axes[0].plot(epochs, history["train_loss"], label="train") axes[0].plot(epochs, history["val_loss"], label="val") axes[0].set_yscale("log") axes[0].set_xlabel("epoch") axes[0].set_ylabel("total loss") axes[0].legend() axes[0].grid(alpha=.3) axes[1].plot(epochs, history["efficiency"], label="efficiency") axes[1].plot(epochs, history["purity"], label="purity") axes[1].set_xlabel("epoch") axes[1].set_ylim(0, 1.02) axes[1].legend() axes[1].grid(alpha=.3) axes[2].plot(epochs, history["rms_1d_um"], label="RMS 1D") axes[2].plot(epochs, history["sigma68_1d_um"], label=r"$\sigma_{68}$") axes[2].set_xlabel("epoch") axes[2].set_ylabel(r"residual [$\mu$m]") axes[2].set_yscale("log") axes[2].legend() axes[2].grid(alpha=.3) figure.tight_layout() figure.savefig(output_dir / "Plots" / "training_curves.png", dpi=130) plt.close(figure) def main(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True) parser.add_argument("--epochs", type=int, default=None) parser.add_argument("--tag", default=None) args = parser.parse_args() conf = OmegaConf.load(args.config) if args.epochs: conf.training.epochs = args.epochs torch.manual_seed(conf.training.seed) np.random.seed(conf.training.seed) output_dir = prepare_output_dir(conf, args.tag) print(f"[train] output -> {output_dir}") dataset_kwargs = dict(input_scale=float(conf.data.input_scale), energy_scale=float(conf.data.energy_scale)) train_dataset = GridFrameDataset( shard_files(conf.data.folder, "train", conf.data.train_shards), augment=bool(conf.data.augment), name="train", max_frames=conf.data.max_train_frames, **dataset_kwargs) val_dataset = GridFrameDataset( shard_files(conf.data.folder, "val", conf.data.val_shards), augment=False, name="val", return_photons=True, max_frames=conf.data.max_val_frames, **dataset_kwargs) n_workers = int(conf.data.num_workers) train_loader = DataLoader(train_dataset, batch_size=int(conf.data.batch_size_train), shuffle=True, num_workers=n_workers, pin_memory=True, drop_last=True, persistent_workers=n_workers > 0) val_loader = DataLoader(val_dataset, batch_size=int(conf.data.batch_size_val), shuffle=False, num_workers=n_workers, pin_memory=True, collate_fn=collate_with_photons, persistent_workers=n_workers > 0) model = build_model(conf).cuda() halo = model.rf_radius frame_size = train_dataset.frame_size mask = valid_mask(frame_size, frame_size, halo, device="cuda") n_parameters = sum(p.numel() for p in model.parameters()) print(f"[train] GridNet: {n_parameters} parameters, receptive-field radius " f"R={halo} px (halo), valid region " f"{frame_size-2*halo}x{frame_size-2*halo} of {frame_size}x{frame_size}") loss_function = GridLoss(conf).cuda() optimizer = torch.optim.Adam(model.parameters(), lr=float(conf.training.lr), weight_decay=float(conf.training.weight_decay)) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, "min", factor=float(conf.training.scheduler_factor), patience=int(conf.training.scheduler_patience)) use_amp = bool(conf.training.amp) grad_scaler = torch.amp.GradScaler("cuda") if use_amp else None history = {key: [] for key in ("train_loss", "val_loss", "efficiency", "purity", "rms_1d_um", "sigma68_1d_um")} best_loss = float("inf") for epoch in range(1, int(conf.training.epochs) + 1): started = time.time() train_metrics = run_epoch(model, train_loader, loss_function, mask, optimizer, grad_scaler, use_amp) val_metrics = run_epoch(model, val_loader, loss_function, mask, None, None, use_amp) detection = run_detection_metrics(model, val_loader, conf, halo, frame_size, use_amp) scheduler.step(val_metrics["loss"]) history["train_loss"].append(train_metrics["loss"]) history["val_loss"].append(val_metrics["loss"]) history["efficiency"].append(detection["efficiency"]) history["purity"].append(detection["purity"]) history["rms_1d_um"].append(detection["rms_1d_um"]) history["sigma68_1d_um"].append(detection["sigma68_1d_um"]) print(f"[{epoch:3d}/{conf.training.epochs}] " f"train {train_metrics['loss']:.4f} " f"(obj {train_metrics['objectness']:.4f} " f"off {train_metrics['offset']:.5f} " f"ene {train_metrics['energy']:.5f}) " f"val {val_metrics['loss']:.4f} | " f"eff {detection['efficiency']*100:5.2f}% " f"pur {detection['purity']*100:5.2f}% " f"| RMS {detection['rms_1d_um']:.3f} um " f"s68 {detection['sigma68_1d_um']:.3f} um " f"| lr {optimizer.param_groups[0]['lr']:.1e} " f"| {time.time()-started:.0f}s", flush=True) checkpoint = dict(model=model.state_dict(), conf=OmegaConf.to_container(conf), epoch=epoch, metrics=detection) if val_metrics["loss"] < best_loss: best_loss = val_metrics["loss"] torch.save(checkpoint, output_dir / "Models" / "best.pth") if epoch in list(conf.training.checkpoint_epochs) or epoch == conf.training.epochs: torch.save(checkpoint, output_dir / "Models" / f"E{epoch}.pth") torch.save(checkpoint, output_dir / "Models" / "last.pth") plot_curves(history, output_dir) with open(output_dir / "history.json", "w") as handle: json.dump(history, handle, indent=1) print(f"[train] done. best val loss {best_loss:.4f}. " f"models in {output_dir/'Models'}") if __name__ == "__main__": main()