"""Train the requested full 500-epoch unlabeled VINO Darcy experiment.""" from __future__ import annotations import csv import time import platform import random import numpy as np import torch import yaml from common import (DarcyEnergy, arguments, build_model, configuration, evaluate, initialize, output_path, read_data, save_checkpoint, write_json) def train(): parser = arguments(__doc__) parser.add_argument("--resume", type=str, help="Restore model, optimizer, scheduler and RNG from last_model.pt") args = parser.parse_args() cfg = configuration(args) if cfg["training"]["epochs"] != 500: raise ValueError("The formal run must complete exactly 500 epochs") device = initialize(cfg) settings = cfg["training"] results = output_path(cfg["output"]["results_dir"]) results.mkdir(parents=True, exist_ok=True) fields, metadata = read_data(cfg) write_json(results / "data_manifest.json", metadata) (results / "resolved_config.yaml").write_text(yaml.safe_dump(cfg, sort_keys=False)) model = build_model(cfg, device) energy = DarcyEnergy(cfg["data"]["resolution"], cfg["physics"]["source"]).to(device) p_train = fields["P_train"].to(device) p_test, u_test = fields["P_test"].to(device), fields["U_test"].to(device) # Only permeability and its stiffness enter optimization. U_train is never put on device. with torch.no_grad(): stiffness = energy.stiffness(p_train) optimizer = torch.optim.Adam(model.parameters(), lr=settings["learning_rate"], betas=tuple(settings["betas"]), eps=settings["eps"], weight_decay=settings["weight_decay"]) scheduler_cfg = dict(settings["scheduler"]) scheduler_cfg.pop("type") scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, **scheduler_cfg) best, best_epoch, start_epoch, previous_seconds = float("inf"), 0, 0, 0.0 history_rows = [] shuffle_rng = torch.Generator().manual_seed(cfg["seed"]) if args.resume: checkpoint = torch.load(args.resume, map_location=device, weights_only=False) if checkpoint["config"] != cfg or checkpoint["data_sha256"] != metadata["sha256"]: raise ValueError("Resume configuration or original data hash changed") model.load_state_dict(checkpoint["model_state_dict"]) optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) start_epoch, best, best_epoch = checkpoint["epoch"], checkpoint["best_relative_l2"], checkpoint["best_epoch"] previous_seconds = checkpoint["elapsed_seconds"] history_rows = checkpoint["history"] test_l2, test_std = checkpoint["test_relative_l2"], checkpoint["test_relative_l2_std"] torch.set_rng_state(checkpoint["torch_rng"].cpu()) shuffle_rng.set_state(checkpoint["shuffle_rng"].cpu()) random.setstate(checkpoint["python_rng"]) np.random.set_state(checkpoint["numpy_rng"]) if device.type == "cuda": torch.cuda.set_rng_state_all([state.cpu() for state in checkpoint["accelerator_rng"]]) elif (results / "history.csv").exists(): raise FileExistsError("Training history exists; use --resume to continue without overwriting it") environment = {"python": platform.python_version(), "torch": str(torch.__version__), "numpy": np.__version__, "hip": torch.version.hip, "cuda": torch.version.cuda, "device": str(device), "device_name": torch.cuda.get_device_name() if device.type == "cuda" else platform.processor(), "parameters_real_equivalent": sum(p.numel() * (2 if p.is_complex() else 1) for p in model.parameters()), "training_labels_used": False, "test_used_for_checkpoint_selection": True} write_json(results / "environment.json", environment) print(f"Environment: {environment}", flush=True) print(f"Original dataset: {metadata['path']} sha256={metadata['sha256']}; 1000 train / 100 test", flush=True) print("Objective: exact Q1 Darcy energy; no solution-label loss; full 500 epochs", flush=True) columns = ["epoch", "train_energy", "test_energy", "test_relative_l2", "test_relative_l2_std", "learning_rate", "epoch_seconds", "elapsed_seconds", "best_epoch"] run_start = time.perf_counter() # Rebuild the CSV from the atomically saved state on resume, including a # final row whose CSV write may have been interrupted after checkpointing. with (results / "history.csv").open("w", newline="", buffering=1) as history: writer = csv.DictWriter(history, fieldnames=columns) writer.writeheader() writer.writerows(history_rows) for epoch in range(start_epoch + 1, settings["epochs"] + 1): epoch_start = time.perf_counter() model.train() permutation = torch.randperm(len(p_train), generator=shuffle_rng).to(device) total_energy = 0.0 for indices in permutation.split(settings["batch_size"]): optimizer.zero_grad(set_to_none=True) predicted = model(p_train[indices]) loss = energy(predicted, stiffness=stiffness[indices]) if not torch.isfinite(loss): raise FloatingPointError(f"Nonfinite training energy at epoch {epoch}") loss.backward() optimizer.step() total_energy += loss.item() * len(indices) train_energy = total_energy / len(p_train) # Match the official sum of mean batch energies, independent of test labels. scheduler.step(total_energy / settings["batch_size"]) _, errors, test_energies = evaluate(model, p_test, u_test, energy, settings["batch_size"]) test_l2 = errors.double().mean().item() test_std = errors.double().std(correction=1).item() if not torch.isfinite(errors).all(): raise FloatingPointError(f"Nonfinite test metric at epoch {epoch}") improved = test_l2 < best if improved: best, best_epoch = test_l2, epoch elapsed = previous_seconds + time.perf_counter() - run_start row = dict(epoch=epoch, train_energy=train_energy, test_energy=test_energies.double().mean().item(), test_relative_l2=test_l2, test_relative_l2_std=test_std, learning_rate=optimizer.param_groups[0]["lr"], epoch_seconds=time.perf_counter() - epoch_start, elapsed_seconds=elapsed, best_epoch=best_epoch) history_rows.append(row) state = {"model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "scheduler_state_dict": scheduler.state_dict(), "epoch": epoch, "config": cfg, "best_relative_l2": best, "best_epoch": best_epoch, "test_relative_l2": test_l2, "test_relative_l2_std": test_std, "train_energy": train_energy, "elapsed_seconds": elapsed, "data_sha256": metadata["sha256"], "torch_rng": torch.get_rng_state(), "shuffle_rng": shuffle_rng.get_state(), "python_rng": random.getstate(), "numpy_rng": np.random.get_state(), "accelerator_rng": torch.cuda.get_rng_state_all() if device.type == "cuda" else [], "history": history_rows} if improved: save_checkpoint(output_path(cfg["output"]["checkpoint"]), state) save_checkpoint(output_path(cfg["output"]["last_checkpoint"]), state) writer.writerow(row) write_json(results / "progress.json", {**row, "target_epochs": 500, "best_relative_l2": best, "status": "completed" if epoch == 500 else "running"}) print(f"Epoch {epoch:03d}/500 loss(energy)={train_energy:.9e} test_energy={row['test_energy']:.9e} " f"test_rel_L2={test_l2:.8f} ({100*test_l2:.4f}%) std={100*test_std:.4f}% " f"lr={row['learning_rate']:.3e} best={100*best:.4f}%@{best_epoch} time={row['epoch_seconds']:.2f}s", flush=True) write_json(results / "training_summary.json", {"status": "completed", "completed_epochs": 500, "best_epoch": best_epoch, "best_test_relative_l2": best, "last_test_relative_l2": test_l2, "last_test_relative_l2_std": test_std, "training_seconds": previous_seconds + time.perf_counter() - run_start, "early_stopping": False, "data_sha256": metadata["sha256"], "training_labels_used": False}) print(f"COMPLETE: 500 epochs; best test relative L2={best:.8f} at epoch {best_epoch}", flush=True) if __name__ == "__main__": train()