| """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() |
| |
| |
| 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() |
|
|
