File size: 6,022 Bytes
8bf8ea0 | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | from __future__ import annotations
import csv
import json
import os
import random
import sys
from pathlib import Path
from typing import Any, Dict, Tuple
import h5py
import numpy as np
import torch
import yaml
from torch.utils.data import DataLoader, TensorDataset
PACKAGE_ROOT = Path(__file__).resolve().parents[1]
if str(PACKAGE_ROOT) not in sys.path:
sys.path.insert(0, str(PACKAGE_ROOT))
from model.crop_navier_stokes import CROPNavierStokes, count_parameters
def load_config(path: str | os.PathLike[str]) -> Dict[str, Any]:
with open(path, "r", encoding="utf-8") as f:
cfg = yaml.safe_load(f)
cfg["_config_path"] = str(Path(path).resolve())
cfg["_package_root"] = str(PACKAGE_ROOT)
return cfg
def resolve_output_path(cfg: Dict[str, Any], relative: str) -> Path:
out = Path(cfg["paths"]["output_dir"])
path = Path(relative)
return path if path.is_absolute() else out / path
def set_seed(seed: int, deterministic: bool = False, benchmark: bool = True) -> None:
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = deterministic
torch.backends.cudnn.benchmark = benchmark
def get_device(cfg: Dict[str, Any]) -> torch.device:
requested = cfg.get("runtime", {}).get("device", "cuda")
if requested == "cuda" and torch.cuda.is_available():
return torch.device("cuda")
return torch.device("cpu")
def build_model(cfg: Dict[str, Any]) -> CROPNavierStokes:
m = cfg["model"]
return CROPNavierStokes(
in_channels=cfg["data"]["channels"],
out_channels=cfg["data"]["channels"],
modes=m["modes"],
initial_channels=m["initial_channels"],
levels=m["levels"],
residual_blocks_per_level=m["residual_blocks_per_level"],
bottleneck_blocks=m["bottleneck_blocks"],
kernel_size=m["kernel_size"],
input_size=m["input_size"],
latent_size=m["latent_size"],
use_bn=m["use_bn"],
)
def load_tensors(cfg: Dict[str, Any]) -> Tuple[torch.Tensor, torch.Tensor]:
data_path = Path(cfg["paths"]["data_dir"]) / cfg["paths"]["data_file"]
key = cfg["data"]["key"]
input_t = cfg["data"]["input_time_index"]
target_t = cfg["data"]["target_time_index"]
with h5py.File(data_path, "r") as f:
if key not in f:
raise KeyError(f"HDF5 key {key!r} not found in {data_path}")
dset = f[key]
expected = tuple(cfg["data"]["raw_shape"])
if tuple(dset.shape) != expected:
raise ValueError(f"Expected {key} shape {expected}, got {tuple(dset.shape)}")
x_np = np.asarray(dset[input_t, :, :, :], dtype=np.float32).transpose(2, 0, 1)
y_np = np.asarray(dset[target_t, :, :, :], dtype=np.float32).transpose(2, 0, 1)
x = torch.from_numpy(x_np[:, None, :, :]).contiguous()
y = torch.from_numpy(y_np[:, None, :, :]).contiguous()
return x, y
def make_loaders(cfg: Dict[str, Any]) -> Tuple[DataLoader, DataLoader, DataLoader]:
x, y = load_tensors(cfg)
n_train = cfg["data"]["train_size"]
n_val = cfg["data"]["val_size"]
n_test = cfg["data"]["test_size"]
n_total = n_train + n_val + n_test
if x.shape[0] < n_total:
raise ValueError(f"Need {n_total} samples, found {x.shape[0]}")
train = TensorDataset(x[:n_train], y[:n_train])
val = TensorDataset(x[n_train : n_train + n_val], y[n_train : n_train + n_val])
test = TensorDataset(x[n_train + n_val : n_total], y[n_train + n_val : n_total])
batch_size = cfg["training"]["batch_size"]
workers = cfg["training"].get("num_workers", 0)
return (
DataLoader(train, batch_size=batch_size, shuffle=True, num_workers=workers, pin_memory=torch.cuda.is_available()),
DataLoader(val, batch_size=batch_size, shuffle=False, num_workers=workers, pin_memory=torch.cuda.is_available()),
DataLoader(test, batch_size=batch_size, shuffle=False, num_workers=workers, pin_memory=torch.cuda.is_available()),
)
def relative_normalized_mse(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
return torch.mean((pred - target) ** 2) / torch.mean(target**2).clamp_min(1e-12)
@torch.no_grad()
def mean_relative_l2_percent(model: torch.nn.Module, loader: DataLoader, device: torch.device) -> float:
model.eval()
total = 0.0
count = 0
for x, y in loader:
x = x.to(device, non_blocking=True)
y = y.to(device, non_blocking=True)
pred = model(x)
diff = torch.linalg.vector_norm((pred - y).reshape(pred.shape[0], -1), ord=2, dim=1)
denom = torch.linalg.vector_norm(y.reshape(y.shape[0], -1), ord=2, dim=1).clamp_min(1e-12)
rel = diff / denom
total += rel.sum().item()
count += rel.numel()
return 100.0 * total / max(1, count)
@torch.no_grad()
def collect_predictions(model: torch.nn.Module, loader: DataLoader, device: torch.device) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
model.eval()
preds, targets, inputs = [], [], []
for x, y in loader:
pred = model(x.to(device, non_blocking=True)).cpu()
preds.append(pred.numpy().astype(np.float32))
targets.append(y.numpy().astype(np.float32))
inputs.append(x.numpy().astype(np.float32))
return np.concatenate(inputs), np.concatenate(preds), np.concatenate(targets)
def save_json(path: Path, payload: Dict[str, Any]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
tmp = path.with_suffix(path.suffix + ".tmp")
with open(tmp, "w", encoding="utf-8") as f:
json.dump(payload, f, indent=2, sort_keys=True)
tmp.replace(path)
def append_history(path: Path, row: Dict[str, Any]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
exists = path.exists()
with open(path, "a", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(f, fieldnames=list(row.keys()))
if not exists:
writer.writeheader()
writer.writerow(row)
|