"""Formal 300-epoch FluxNet-LAP pushforward training on SCNet DCU.""" import csv import json import random import time from pathlib import Path import numpy as np import torch from torch.nn import functional as F from torch.utils.data import DataLoader from common import (ROOT, Trajectories, Windows, build_model, checkpoint_config_matches, device, load_config, output_dir, parser, provenance, rollout_metrics, save_json, setup, validate) def train_batch(model, window, optimizer, steps): """Eq. (6): first-step MSE + terminal MSE; every feedback is detached.""" optimizer.zero_grad(set_to_none=True) prediction = model(window[:, 0]) first_loss = F.mse_loss(prediction, window[:, 1]) current = prediction.detach() # Free the first graph before the terminal graph to reduce memory without # changing gradients or updating parameters between the two losses. first_loss.backward() with torch.no_grad(): for _ in range(1, steps - 1): current = model(current).detach() prediction = model(current.detach()) terminal_loss = F.mse_loss(prediction, window[:, steps]) terminal_loss.backward() if not torch.isfinite(first_loss + terminal_loss): raise FloatingPointError("Nonfinite training loss") if not torch.stack([torch.isfinite(parameter.grad).all() for parameter in model.parameters() if parameter.grad is not None]).all(): raise FloatingPointError("Nonfinite parameter gradient") optimizer.step() return float(first_loss.detach()), float(terminal_loss.detach()) def atomic_checkpoint(path, payload): path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(".tmp") torch.save(payload, temporary) temporary.replace(path) def main(): args_parser = parser(__doc__) args_parser.add_argument("--resume", type=Path) args = args_parser.parse_args() cfg = load_config(args.config) setup(cfg["seed"]) target_device = device() out = output_dir(cfg) settings = cfg["training"] run_info = provenance() print("FORMAL_TRAINING " + json.dumps({**run_info, "config": cfg}), flush=True) train = Trajectories(cfg, "train") val = Trajectories(cfg, "val") save_json(out / "data_manifest.json", {"train": train.metadata, "val": val.metadata}) generator = torch.Generator().manual_seed(cfg["seed"]) train_loader = DataLoader(Windows(train, settings["pushforward_steps"]), batch_size=settings["batch_size"], shuffle=True, num_workers=settings["num_workers"], pin_memory=True, persistent_workers=settings["num_workers"] > 0, generator=generator, drop_last=False) val_loader = DataLoader(Windows(val, settings["pushforward_steps"]), batch_size=cfg["evaluation"]["batch_size"], shuffle=False, num_workers=0, pin_memory=True) model = build_model(cfg, target_device) optimizer = torch.optim.AdamW(model.parameters(), lr=settings["learning_rate"], weight_decay=settings["weight_decay"]) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, factor=settings["scheduler_factor"], patience=settings["scheduler_patience"]) epoch_start, best_loss, best_epoch, history = 1, float("inf"), 0, [] if args.resume: checkpoint = torch.load(args.resume, map_location="cpu", weights_only=False) if not checkpoint_config_matches(checkpoint["config"], cfg): raise ValueError("Resume config differs from checkpoint") model.load_state_dict(checkpoint["model_state_dict"], strict=True) optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) epoch_start = checkpoint["epoch"] + 1 best_loss, best_epoch = checkpoint["best_val_loss"], checkpoint["best_epoch"] history = checkpoint["history"] random.setstate(checkpoint["random_state"]) np.random.set_state(checkpoint["numpy_state"]) torch.set_rng_state(checkpoint["torch_state"]) torch.cuda.set_rng_state_all(checkpoint["cuda_state"]) generator.set_state(checkpoint["loader_state"]) print(f"TRAIN samples={len(train_loader.dataset)} batches/epoch={len(train_loader)} " f"parameters={sum(p.numel() for p in model.parameters())} " f"epochs={settings['epochs']} start_epoch={epoch_start}", flush=True) started = time.monotonic() long_data = None for epoch in range(epoch_start, settings["epochs"] + 1): before = time.monotonic() model.train() losses = np.zeros(2) samples = 0 lr = optimizer.param_groups[0]["lr"] for batch_index, window in enumerate(train_loader, 1): window = window.to(target_device, non_blocking=True) first, terminal = train_batch(model, window, optimizer, settings["pushforward_steps"]) losses += np.array([first, terminal]) * len(window) samples += len(window) if batch_index % settings["print_every_batches"] == 0 or batch_index == len(train_loader): print(f"epoch={epoch:03d}/300 batch={batch_index:03d}/{len(train_loader)} " f"loss={(losses.sum()/samples):.8e} one_step={losses[0]/samples:.8e} " f"pushforward5={losses[1]/samples:.8e} lr={lr:.3e}", flush=True) validation = validate(model, val_loader, target_device) scheduler.step(validation["mse"]) if validation["mse"] < best_loss: best_loss, best_epoch = validation["mse"], epoch row = {"epoch": epoch, "train_loss": float(losses.sum()/samples), "train_step1_mse": float(losses[0]/samples), "train_step5_mse": float(losses[1]/samples), "val_mse": validation["mse"], "val_step1_mse": validation["step1_mse"], "val_step5_mse": validation["step5_mse"], **{f"val_mae_{name}": x for name, x in zip(("h", "mx", "my"), validation["mae"])}, "val_negative_h_ratio": validation["negative_h_ratio"], "lr": lr, "seconds": time.monotonic() - before} history.append(row) print("EPOCH " + json.dumps(row), flush=True) save_json(out / "training_history.json", history) with open(out / "training_history.csv", "w", newline="") as stream: writer = csv.DictWriter(stream, fieldnames=row.keys()) writer.writeheader() writer.writerows(history) payload = {"model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "epoch": epoch, "best_val_loss": best_loss, "best_epoch": best_epoch, "config": cfg, "history": history, "provenance": run_info, "random_state": random.getstate(), "numpy_state": np.random.get_state(), "torch_state": torch.get_rng_state(), "cuda_state": torch.cuda.get_rng_state_all(), "loader_state": generator.get_state()} if best_epoch == epoch: atomic_checkpoint(ROOT / cfg["paths"]["best_checkpoint"], payload) print(f"BEST checkpoint epoch={best_epoch} val_mse={best_loss:.8e}", flush=True) if epoch % settings["checkpoint_every_epochs"] == 0 or epoch == settings["epochs"]: atomic_checkpoint(ROOT / cfg["paths"]["last_checkpoint"], payload) save_json(out / "training_status.json", {"state": "running", "completed_epochs": epoch, "best_epoch": best_epoch, "best_val_loss": best_loss, "elapsed_seconds": time.monotonic() - started, "provenance": run_info}) if epoch % cfg["evaluation"]["long_test_every_epochs"] == 0: if long_data is None: long_data = Trajectories(cfg, "test_long") diagnostic = rollout_metrics(model, long_data, target_device, cfg["evaluation"]["batch_size"]) save_json(out / f"diagnostic_long_epoch{epoch:03d}.json", diagnostic) print("LONG_TEST_DIAGNOSTIC " + json.dumps({"epoch": epoch, **diagnostic}), flush=True) assert len(history) == settings["epochs"] and history[-1]["epoch"] == 300 save_json(out / "training_status.json", {"state": "completed", "completed_epochs": 300, "best_epoch": best_epoch, "best_val_loss": best_loss, "elapsed_seconds": time.monotonic() - started, "provenance": run_info, "formal_training": True, "stopped_early": False}) print(f"TRAINING_COMPLETED epochs=300 best_epoch={best_epoch} best_val_mse={best_loss:.8e}", flush=True) if __name__ == "__main__": main()