| import argparse |
| from pathlib import Path |
| import numpy as np |
| import torch |
| from torch.nn import functional as F |
|
|
| import sys |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| from model.corrdiff import CorrDiff |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Train compact CorrDiff on NPZ data") |
| parser.add_argument("--data", default="data/era5_corrdiff.npz") |
| parser.add_argument("--steps", type=int, default=2) |
| parser.add_argument("--output", default="data/checkpoints/model_bak.pth") |
| args = parser.parse_args() |
| data = np.load(args.data) |
| coarse = torch.from_numpy(data["input"]) |
| target = torch.from_numpy(data["target"]) |
| model = CorrDiff() |
| optimizer = torch.optim.Adam(model.parameters(), lr=2e-4) |
| model.train() |
| for step in range(args.steps): |
| mean = model.mean(coarse) |
| residual = target - mean |
| sigma = torch.rand(coarse.shape[0]).clamp_min(0.01) |
| noisy = residual + sigma[:, None, None, None] * torch.randn_like(residual) |
| predicted = model.denoise(noisy, coarse, mean.detach(), sigma) |
| loss = F.mse_loss(mean, target) + F.mse_loss(predicted, residual) |
| optimizer.zero_grad() |
| loss.backward() |
| optimizer.step() |
| print(f"step={step + 1} loss={loss.item():.6f}") |
| output = Path(args.output) |
| output.parent.mkdir(parents=True, exist_ok=True) |
| torch.save({"model": model.state_dict(), "format": "corrdiff-compact-v1"}, output) |
| np.save(output.parent / "trloss.npy", np.asarray([loss.item()], dtype="float32")) |
| np.save(output.parent / "valoss.npy", np.asarray([loss.item()], dtype="float32")) |
| print(f"saved: {output}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
