| from __future__ import annotations |
|
|
| import argparse |
| import copy |
| from pathlib import Path |
|
|
| import torch |
|
|
| from common import ( |
| append_history, |
| build_model, |
| collect_predictions, |
| count_parameters, |
| get_device, |
| load_config, |
| make_loaders, |
| mean_relative_l2_percent, |
| relative_normalized_mse, |
| resolve_output_path, |
| save_json, |
| set_seed, |
| ) |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", default=str(Path(__file__).resolve().parents[1] / "config" / "config.yaml")) |
| parser.add_argument("--dry-run", action="store_true", help="Run one forward/backward pass for preflight only.") |
| return parser.parse_args() |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| cfg = load_config(args.config) |
| set_seed(cfg["training"]["seed"], cfg["runtime"]["deterministic"], cfg["runtime"]["benchmark"]) |
| device = get_device(cfg) |
|
|
| output_dir = Path(cfg["paths"]["output_dir"]) |
| weight_path = resolve_output_path(cfg, cfg["paths"]["checkpoint"]) |
| results_dir = resolve_output_path(cfg, cfg["paths"]["results_dir"]) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| weight_path.parent.mkdir(parents=True, exist_ok=True) |
| results_dir.mkdir(parents=True, exist_ok=True) |
|
|
| train_loader, val_loader, test_loader = make_loaders(cfg) |
| model = build_model(cfg).to(device) |
| params = count_parameters(model) |
|
|
| optimizer = torch.optim.AdamW( |
| model.parameters(), |
| lr=cfg["training"]["learning_rate"], |
| weight_decay=cfg["training"]["weight_decay"], |
| ) |
| scheduler = torch.optim.lr_scheduler.StepLR( |
| optimizer, |
| step_size=cfg["training"]["scheduler_step_size"], |
| gamma=cfg["training"]["scheduler_gamma"], |
| ) |
|
|
| if args.dry_run: |
| model.train() |
| x, y = next(iter(train_loader)) |
| x = x.to(device) |
| y = y.to(device) |
| loss = relative_normalized_mse(model(x), y) |
| loss.backward() |
| print(f"DRY_RUN forward/backward passed loss={loss.item():.8f} params={params}", flush=True) |
| return |
|
|
| best_val = float("inf") |
| best_test = float("inf") |
| best_epoch = -1 |
| best_state = None |
| patience = cfg["training"]["early_stopping_patience"] |
| stale_epochs = 0 |
| history_path = results_dir / "history.csv" |
|
|
| for epoch in range(1, cfg["training"]["epochs"] + 1): |
| model.train() |
| running_loss = 0.0 |
| batches = 0 |
| for x, y in train_loader: |
| x = x.to(device, non_blocking=True) |
| y = y.to(device, non_blocking=True) |
| optimizer.zero_grad(set_to_none=True) |
| pred = model(x) |
| loss = relative_normalized_mse(pred, y) |
| loss.backward() |
| optimizer.step() |
| running_loss += loss.item() |
| batches += 1 |
|
|
| train_loss = running_loss / max(1, batches) |
| val_l2 = mean_relative_l2_percent(model, val_loader, device) |
| test_l2 = mean_relative_l2_percent(model, test_loader, device) |
| lr = optimizer.param_groups[0]["lr"] |
| row = { |
| "epoch": epoch, |
| "train_relative_normalized_mse": train_loss, |
| "val_mean_relative_l2_percent": val_l2, |
| "test_mean_relative_l2_percent": test_l2, |
| "lr": lr, |
| } |
| append_history(history_path, row) |
| print( |
| f"epoch={epoch:04d} train_loss={train_loss:.8f} " |
| f"val_mean_relative_l2={val_l2:.6f}% test_mean_relative_l2={test_l2:.6f}% lr={lr:.8g}", |
| flush=True, |
| ) |
|
|
| if val_l2 < best_val: |
| best_val = val_l2 |
| best_test = test_l2 |
| best_epoch = epoch |
| best_state = copy.deepcopy(model.state_dict()) |
| torch.save( |
| { |
| "model_state": best_state, |
| "config": cfg, |
| "epoch": best_epoch, |
| "best_val_mean_relative_l2_percent": best_val, |
| "test_mean_relative_l2_percent_at_best_val": best_test, |
| "parameter_count": params, |
| }, |
| weight_path, |
| ) |
| stale_epochs = 0 |
| else: |
| stale_epochs += 1 |
|
|
| save_json( |
| results_dir / "metrics.json", |
| { |
| "status": "running", |
| "epoch": epoch, |
| "best_epoch": best_epoch, |
| "best_val_mean_relative_l2_percent": best_val, |
| "test_mean_relative_l2_percent_at_best_val": best_test, |
| "paper_crop_re10000_percent": cfg["experiment"]["paper_reference_test_percent"], |
| "paper_crop_re10000_std_percent": cfg["experiment"]["paper_reference_std_percent"], |
| "parameter_count": params, |
| }, |
| ) |
| scheduler.step() |
|
|
| if stale_epochs >= patience: |
| print( |
| f"early_stopping epoch={epoch} best_epoch={best_epoch} " |
| f"best_val={best_val:.6f}% best_test={best_test:.6f}%", |
| flush=True, |
| ) |
| break |
|
|
| if best_state is not None: |
| model.load_state_dict(best_state) |
| inputs, preds, targets = collect_predictions(model, test_loader, device) |
| pred_path = results_dir / "test_predictions_from_train.npz" |
| import numpy as np |
|
|
| np.savez_compressed(pred_path, inputs=inputs, predictions=preds, targets=targets) |
| save_json( |
| results_dir / "metrics.json", |
| { |
| "status": "complete", |
| "best_epoch": best_epoch, |
| "best_val_mean_relative_l2_percent": best_val, |
| "test_mean_relative_l2_percent_at_best_val": best_test, |
| "paper_crop_re10000_percent": cfg["experiment"]["paper_reference_test_percent"], |
| "paper_crop_re10000_std_percent": cfg["experiment"]["paper_reference_std_percent"], |
| "parameter_count": params, |
| "checkpoint": str(weight_path), |
| "prediction_file": str(pred_path), |
| }, |
| ) |
| print( |
| f"training_complete best_epoch={best_epoch} best_val={best_val:.6f}% " |
| f"test_at_best_val={best_test:.6f}% checkpoint={weight_path}", |
| flush=True, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
