| import argparse |
| import os |
| import sys |
| import time |
| from datetime import datetime |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(PROJECT_ROOT)) |
|
|
| from scripts.common import ( |
| PROBLEMS, |
| build_models, |
| dump_json, |
| ensure_onescience_path, |
| load_config, |
| resolve_path, |
| save_checkpoint, |
| select_device, |
| ) |
|
|
| ensure_onescience_path() |
| from scripts.topology_optimization import find_TO |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser(description="Train GP_for_TO topology optimization models.") |
| parser.add_argument("--problem", choices=PROBLEMS, default=None) |
| parser.add_argument("--gpu", type=int, default=None) |
| parser.add_argument("--device", default=None) |
| parser.add_argument("--num-iter", type=int, default=None) |
| parser.add_argument("--n-col-domain", type=int, default=None) |
| parser.add_argument("--n-train-per-bc", type=int, default=None) |
| parser.add_argument("--diff-method", choices=("Numerical", "Autograd"), default=None) |
| parser.add_argument("--lr", type=float, default=None) |
| parser.add_argument("--checkpoint-path", default=None) |
| parser.add_argument("--output-dir", default=None) |
| parser.add_argument("--no-plot", action="store_true") |
| return parser.parse_args() |
|
|
|
|
| def main(): |
| args = parse_args() |
| cfg = load_config() |
|
|
| if args.problem: |
| cfg["problem"] = args.problem |
| if args.gpu is not None: |
| cfg["training"]["gpu"] = args.gpu |
| if args.device: |
| cfg["training"]["device"] = args.device |
| if args.num_iter is not None: |
| cfg["training"]["num_iter"] = args.num_iter |
| if args.n_col_domain is not None: |
| cfg["data"]["n_col_domain"] = args.n_col_domain |
| if args.n_train_per_bc is not None: |
| cfg["data"]["n_train_per_bc"] = args.n_train_per_bc |
| if args.diff_method: |
| cfg["training"]["diff_method"] = args.diff_method |
| if args.lr is not None: |
| cfg["training"]["lr_default"] = args.lr |
| if args.checkpoint_path: |
| cfg["training"]["checkpoint_path"] = args.checkpoint_path |
| if args.output_dir: |
| cfg["training"]["output_dir"] = args.output_dir |
| if args.no_plot: |
| cfg["training"]["plot_outputs"] = False |
|
|
| os.chdir(PROJECT_ROOT) |
| device = select_device(cfg["training"]) |
| models, metadata = build_models( |
| cfg, |
| device, |
| n_col_domain=cfg["data"]["n_col_domain"], |
| n_train_per_bc=cfg["data"]["n_train_per_bc"], |
| problem=cfg["problem"], |
| ) |
|
|
| title = f"seed{cfg['seed']}_{cfg['problem']}_{datetime.now().strftime('%B%d_%H-%M')}" |
| start_time = time.time() |
| loss_history = find_TO( |
| model_list=models, |
| num_iter=int(cfg["training"]["num_iter"]), |
| lr_default=float(cfg["training"]["lr_default"]), |
| title=title, |
| problem=cfg["problem"], |
| diff_method=cfg["training"]["diff_method"], |
| checkpoint_steps=list(cfg["training"].get("checkpoints", [])), |
| plot_outputs=bool(cfg["training"].get("plot_outputs", True)), |
| ) |
| elapsed = time.time() - start_time |
|
|
| output_dir = resolve_path(cfg["training"]["output_dir"]) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| np.save(output_dir / "loss_history.npy", np.asarray(loss_history, dtype=np.float64)) |
| summary = { |
| **metadata, |
| "device": str(device), |
| "num_iter": int(cfg["training"]["num_iter"]), |
| "diff_method": cfg["training"]["diff_method"], |
| "elapsed_seconds": elapsed, |
| "final_loss": float(loss_history[-1]) if loss_history else None, |
| } |
| dump_json(summary, output_dir / "training_summary.json") |
| ckpt_path = save_checkpoint(cfg["training"]["checkpoint_path"], models, cfg, summary, loss_history) |
|
|
| print(f"Training finished in {elapsed:.2f}s") |
| print(f"Loss history: {output_dir / 'loss_history.npy'}") |
| print(f"Checkpoint: {ckpt_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
