| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import time |
| from pathlib import Path |
|
|
| import torch |
|
|
| from common import ( |
| ROOT, |
| build_model, |
| count_parameters, |
| load_config, |
| make_normalized_loaders, |
| relative_l2_per_sample, |
| relative_l2_sum, |
| resolve, |
| save_json, |
| select_device, |
| set_seed, |
| ) |
|
|
|
|
| def evaluate(model, loader, y_normalizer, device): |
| model.eval() |
| total = 0.0 |
| count = 0 |
| with torch.no_grad(): |
| for x, y_norm in loader: |
| x = x.to(device, non_blocking=True) |
| y_norm = y_norm.to(device, non_blocking=True) |
| pred = y_normalizer.decode(model(x)) |
| target = y_normalizer.decode(y_norm) |
| rel = relative_l2_per_sample(pred, target) |
| total += float(rel.sum().item()) |
| count += int(rel.numel()) |
| return total / max(1, count) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Train HO-FNO on Pipe horizontal velocity.") |
| parser.add_argument("--config", default=str(ROOT / "config" / "config.yaml")) |
| args = parser.parse_args() |
|
|
| cfg = load_config(args.config) |
| set_seed(int(cfg["experiment"].get("seed", 2026))) |
| output_dir = resolve(cfg["paths"]["output_dir"]) |
| weight_path = resolve(cfg["paths"]["checkpoint"], output_dir) |
| results_dir = resolve(cfg["paths"]["results_dir"], output_dir) |
| weight_path.parent.mkdir(parents=True, exist_ok=True) |
| results_dir.mkdir(parents=True, exist_ok=True) |
|
|
| require_accelerator = bool(cfg["runtime"].get("require_accelerator", False)) |
| device = select_device(require_accelerator=require_accelerator) |
| print(f"device={device}", flush=True) |
|
|
| train_loader, test_loader, x_normalizer, y_normalizer = make_normalized_loaders(cfg) |
| x_normalizer.to(device) |
| y_normalizer.to(device) |
|
|
| model = build_model(cfg).to(device) |
| print(f"parameters={count_parameters(model)}", flush=True) |
|
|
| train_cfg = cfg["training"] |
| optimizer = torch.optim.AdamW( |
| model.parameters(), |
| lr=float(train_cfg["learning_rate"]), |
| weight_decay=float(train_cfg.get("weight_decay", 0.0)), |
| ) |
| scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( |
| optimizer, |
| T_max=int(train_cfg["epochs"]), |
| eta_min=float(train_cfg["final_learning_rate"]), |
| ) |
|
|
| history_path = results_dir / "train_history.csv" |
| best_metric = float("inf") |
| best_epoch = -1 |
| train_size = int(cfg["data"]["train_size"]) |
| epochs = int(train_cfg["epochs"]) |
| start_time = time.time() |
|
|
| with history_path.open("w", newline="", encoding="utf-8") as handle: |
| writer = csv.DictWriter(handle, fieldnames=["epoch", "train_loss", "test_mean_relative_l2", "lr", "seconds"]) |
| writer.writeheader() |
| for epoch in range(1, epochs + 1): |
| model.train() |
| epoch_start = time.time() |
| train_loss_sum = 0.0 |
| for x, y_norm in train_loader: |
| x = x.to(device, non_blocking=True) |
| y_norm = y_norm.to(device, non_blocking=True) |
| optimizer.zero_grad(set_to_none=True) |
| pred = y_normalizer.decode(model(x)) |
| target = y_normalizer.decode(y_norm) |
| loss = relative_l2_sum(pred, target) |
| loss.backward() |
| optimizer.step() |
| train_loss_sum += float(loss.item()) |
|
|
| scheduler.step() |
| train_loss = train_loss_sum / train_size |
| test_l2 = evaluate(model, test_loader, y_normalizer, device) |
| lr = float(optimizer.param_groups[0]["lr"]) |
| seconds = time.time() - epoch_start |
| writer.writerow( |
| { |
| "epoch": epoch, |
| "train_loss": train_loss, |
| "test_mean_relative_l2": test_l2, |
| "lr": lr, |
| "seconds": seconds, |
| } |
| ) |
| handle.flush() |
| print( |
| f"epoch={epoch:04d} train_loss={train_loss:.8f} " |
| f"test_mean_relative_l2={test_l2:.8f} lr={lr:.8e}", |
| flush=True, |
| ) |
|
|
| if test_l2 < best_metric: |
| best_metric = test_l2 |
| best_epoch = epoch |
| torch.save( |
| { |
| "model_state": model.state_dict(), |
| "config": cfg, |
| "epoch": epoch, |
| "best_metric": best_metric, |
| "x_normalizer": x_normalizer.state_dict(), |
| "y_normalizer": y_normalizer.state_dict(), |
| }, |
| weight_path, |
| ) |
| save_json( |
| results_dir / "best_metrics.json", |
| {"best_epoch": best_epoch, "test_mean_relative_l2": best_metric, "checkpoint": str(weight_path)}, |
| ) |
|
|
| torch.save( |
| { |
| "model_state": model.state_dict(), |
| "config": cfg, |
| "epoch": epochs, |
| "best_metric": best_metric, |
| "x_normalizer": x_normalizer.state_dict(), |
| "y_normalizer": y_normalizer.state_dict(), |
| }, |
| weight_path.parent / "last_model.pth", |
| ) |
| save_json( |
| results_dir / "training_summary.json", |
| { |
| "epochs": epochs, |
| "best_epoch": best_epoch, |
| "best_test_mean_relative_l2": best_metric, |
| "total_seconds": time.time() - start_time, |
| "checkpoint": str(weight_path), |
| }, |
| ) |
| print(f"best_checkpoint={weight_path}", flush=True) |
| print(f"best_epoch={best_epoch} best_test_mean_relative_l2={best_metric:.8f}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
