OneScience-Group/NNCAM

1#!/usr/bin/env python32import argparse3import json4from pathlib import Path5 6import numpy as np7import torch8from torch.utils.data import DataLoader, TensorDataset9 10from model.nncam import build_model, fit_normalizer, normalize_input, scale_output11 12 13ROOT = Path(__file__).resolve().parents[1]14 15 16def main():17    parser = argparse.ArgumentParser(description="Train NNCAM on the prepared NPZ dataset.")18    parser.add_argument("--data", type=Path, default=ROOT / "data/nncam_fake.npz")19    parser.add_argument("--checkpoint", type=Path, default=ROOT / "result/checkpoints/nncam.pt")20    parser.add_argument("--metrics", type=Path, default=ROOT / "result/training/metrics.json")21    parser.add_argument("--epochs", type=int)22    parser.add_argument("--batch-size", type=int)23    parser.add_argument("--width", type=int)24    parser.add_argument("--depth", type=int)25    parser.add_argument("--paper-model", action="store_true", help="Explicitly use depth=9, width=256, epochs=18, batch_size=1024 (567361 parameters).")26    parser.add_argument("--lr", type=float, default=1e-3)27    parser.add_argument("--seed", type=int, default=42)28    args = parser.parse_args()29    if not args.data.is_file():30        raise FileNotFoundError(f"missing dataset {args.data}; run python scripts/fake_data.py first")31    defaults = {"depth": 9, "width": 256, "epochs": 18, "batch_size": 1024} if args.paper_model else {"depth": 4, "width": 32, "epochs": 3, "batch_size": 64}32    depth, width = args.depth or defaults["depth"], args.width or defaults["width"]33    epochs, batch_size = args.epochs or defaults["epochs"], args.batch_size or defaults["batch_size"]34    torch.manual_seed(args.seed)35    with np.load(args.data) as data:36        x, y = data["x"].astype(np.float32), data["y"].astype(np.float32)37    if x.ndim != 2 or x.shape[1] != 94 or y.shape != (x.shape[0], 65):38        raise ValueError(f"expected x=[N,94], y=[N,65], got {x.shape}, {y.shape}")39    input_mean, input_scale = fit_normalizer(x)40    scaled_y = scale_output(y)41    target_mean, target_scale = fit_normalizer(scaled_y)42    dataset = TensorDataset(torch.from_numpy(normalize_input(x, input_mean, input_scale)), torch.from_numpy(normalize_input(scaled_y, target_mean, target_scale)))43    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)44    model = build_model(width=width, depth=depth)45    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)46    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.2)47    history = []48    for epoch in range(epochs):49        total = 0.050        for xb, yb in loader:51            optimizer.zero_grad(set_to_none=True)52            loss = torch.nn.functional.mse_loss(model(xb), yb)53            loss.backward()54            optimizer.step()55            total += loss.item() * len(xb)56        history.append(total / len(dataset))57        scheduler.step()58        print(f"epoch={epoch + 1:02d} loss={history[-1]:.6f}")59    parameter_count = sum(parameter.numel() for parameter in model.parameters())60    checkpoint = {61        "format_version": 1,62        "model": model.state_dict(),63        "model_config": model.model_config,64        "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)},65        "training": {"epochs": epochs, "batch_size": batch_size, "learning_rate": args.lr, "paper_model": args.paper_model, "parameters": parameter_count},66    }67    args.checkpoint.parent.mkdir(parents=True, exist_ok=True)68    args.metrics.parent.mkdir(parents=True, exist_ok=True)69    torch.save(checkpoint, args.checkpoint)70    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")71    print(f"saved {args.checkpoint}; parameters={parameter_count}; paper_model={args.paper_model}")72 73 74if __name__ == "__main__":75    main()76