| |
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import DataLoader, TensorDataset |
|
|
| from model.nncam import build_model, fit_normalizer, normalize_input, scale_output |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Train NNCAM on the prepared NPZ dataset.") |
| parser.add_argument("--data", type=Path, default=ROOT / "data/nncam_fake.npz") |
| parser.add_argument("--checkpoint", type=Path, default=ROOT / "result/checkpoints/nncam.pt") |
| parser.add_argument("--metrics", type=Path, default=ROOT / "result/training/metrics.json") |
| parser.add_argument("--epochs", type=int) |
| parser.add_argument("--batch-size", type=int) |
| parser.add_argument("--width", type=int) |
| parser.add_argument("--depth", type=int) |
| parser.add_argument("--paper-model", action="store_true", help="Explicitly use depth=9, width=256, epochs=18, batch_size=1024 (567361 parameters).") |
| parser.add_argument("--lr", type=float, default=1e-3) |
| parser.add_argument("--seed", type=int, default=42) |
| args = parser.parse_args() |
| if not args.data.is_file(): |
| raise FileNotFoundError(f"missing dataset {args.data}; run python scripts/fake_data.py first") |
| defaults = {"depth": 9, "width": 256, "epochs": 18, "batch_size": 1024} if args.paper_model else {"depth": 4, "width": 32, "epochs": 3, "batch_size": 64} |
| depth, width = args.depth or defaults["depth"], args.width or defaults["width"] |
| epochs, batch_size = args.epochs or defaults["epochs"], args.batch_size or defaults["batch_size"] |
| torch.manual_seed(args.seed) |
| with np.load(args.data) as data: |
| x, y = data["x"].astype(np.float32), data["y"].astype(np.float32) |
| if x.ndim != 2 or x.shape[1] != 94 or y.shape != (x.shape[0], 65): |
| raise ValueError(f"expected x=[N,94], y=[N,65], got {x.shape}, {y.shape}") |
| input_mean, input_scale = fit_normalizer(x) |
| scaled_y = scale_output(y) |
| target_mean, target_scale = fit_normalizer(scaled_y) |
| dataset = TensorDataset(torch.from_numpy(normalize_input(x, input_mean, input_scale)), torch.from_numpy(normalize_input(scaled_y, target_mean, target_scale))) |
| loader = DataLoader(dataset, batch_size=batch_size, shuffle=True) |
| model = build_model(width=width, depth=depth) |
| optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) |
| scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.2) |
| history = [] |
| for epoch in range(epochs): |
| total = 0.0 |
| for xb, yb in loader: |
| optimizer.zero_grad(set_to_none=True) |
| loss = torch.nn.functional.mse_loss(model(xb), yb) |
| loss.backward() |
| optimizer.step() |
| total += loss.item() * len(xb) |
| history.append(total / len(dataset)) |
| scheduler.step() |
| print(f"epoch={epoch + 1:02d} loss={history[-1]:.6f}") |
| parameter_count = sum(parameter.numel() for parameter in model.parameters()) |
| checkpoint = { |
| "format_version": 1, |
| "model": model.state_dict(), |
| "model_config": model.model_config, |
| "normalization": {"input_mean": torch.from_numpy(input_mean), "input_scale": torch.from_numpy(input_scale), "target_mean": torch.from_numpy(target_mean), "target_scale": torch.from_numpy(target_scale)}, |
| "training": {"epochs": epochs, "batch_size": batch_size, "learning_rate": args.lr, "paper_model": args.paper_model, "parameters": parameter_count}, |
| } |
| args.checkpoint.parent.mkdir(parents=True, exist_ok=True) |
| args.metrics.parent.mkdir(parents=True, exist_ok=True) |
| torch.save(checkpoint, args.checkpoint) |
| args.metrics.write_text(json.dumps({"loss": history, "final_loss": history[-1], "parameters": parameter_count, "model_config": model.model_config}, indent=2), encoding="utf-8") |
| print(f"saved {args.checkpoint}; parameters={parameter_count}; paper_model={args.paper_model}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
