File size: 1,716 Bytes
986404c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | 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()
|