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