File size: 7,474 Bytes
80cf062 | 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 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | from __future__ import annotations
import argparse
from pathlib import Path
import sys
from typing import Any
import numpy as np
import torch
from torch.utils.data import DataLoader
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
SCRIPTS_DIR = PROJECT_ROOT / "scripts"
if str(SCRIPTS_DIR) not in sys.path:
sys.path.insert(0, str(SCRIPTS_DIR))
from model.checkpoint import load_checkpoint
from era5_adapter import WMAEERA5Dataset
from train import build_model, device_summary, load_config, resolve_device, resolve_project_path, validate_config
def build_inference_dataset(config: dict[str, Any], config_path: Path) -> WMAEERA5Dataset:
data = config["data"]
if int(data["input_steps"]) != 1 or int(data["output_steps"]) != 1:
raise ValueError("W-MAE inference currently requires input_steps=1 and output_steps=1.")
if not data["variables"]:
raise ValueError(
"data.variables is empty. Supply the verified, explicitly ordered 20-channel list before inference."
)
return WMAEERA5Dataset(
dataset_dir=resolve_project_path(data["dataset_dir"], config_path),
years=data["test_years"],
variables=data["variables"],
task="pretrain",
input_steps=1,
output_steps=1,
normalize=bool(data["normalize"]),
)
def sample_time_index(time_index: Any, batch_index: int, batch_size: int) -> Any:
"""Extract one sample from DataLoader's sequence-major default collation."""
if isinstance(time_index, (list, tuple)):
if time_index and isinstance(time_index[0], (list, tuple)):
return time_index[0][batch_index]
if len(time_index) == batch_size:
return time_index[batch_index]
return time_index
def run_inference(
model: torch.nn.Module,
loader: DataLoader,
output_dir: Path,
mask_ratio: float,
device: torch.device,
max_samples: int | None = None,
log_interval: int = 1,
) -> int:
output_dir.mkdir(parents=True, exist_ok=True)
model.eval()
written = 0
with torch.no_grad():
total_batches = len(loader)
for batch_index, (inputs, _, _, step_idx, time_index) in enumerate(loader, start=1):
if inputs.ndim != 4:
raise ValueError(f"W-MAE inference expects a 4D batch, got {tuple(inputs.shape)}.")
inputs = inputs.to(device)
result = model(inputs, mask_ratio=mask_ratio)
reconstruction = model.unpatchify(result.prediction)
error = reconstruction - inputs
mask = result.mask.reshape(result.mask.shape[0], *model.patch_embed.grid_size)
for sample_index in range(inputs.shape[0]):
if max_samples is not None and written >= max_samples:
print(
f"inference batch={batch_index}/{total_batches} samples_written={written}",
flush=True,
)
return written
sample_time = sample_time_index(time_index, sample_index, inputs.shape[0])
if isinstance(sample_time, (list, tuple)):
sample_time = sample_time[0]
np.savez_compressed(
output_dir / f"sample_{written:06d}.npz",
input=inputs[sample_index].cpu().numpy(),
reconstruction=reconstruction[sample_index].cpu().numpy(),
error=error[sample_index].cpu().numpy(),
mask=mask[sample_index].cpu().numpy(),
step_idx=np.asarray(step_idx[sample_index].item()),
time_index=np.asarray(sample_time),
)
written += 1
if batch_index == 1 or batch_index % log_interval == 0 or batch_index == total_batches:
print(
f"inference batch={batch_index}/{total_batches} samples_written={written}",
flush=True,
)
return written
def main() -> None:
parser = argparse.ArgumentParser(description="Run W-MAE reconstruction inference on ERA5 samples.")
parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "conf" / "config.yaml")
parser.add_argument("--checkpoint", type=Path, default="./data/checkpoint/model_bak.pth")
parser.add_argument("--checkpoint-source-root", type=Path, default=None)
parser.add_argument("--non-strict-checkpoint", action="store_true")
parser.add_argument("--output-dir", type=Path, default=None)
parser.add_argument(
"--device",
default="auto",
help="Inference device: auto (default), cuda (AMD DCU through HIP), or cpu.",
)
parser.add_argument("--mask-ratio", type=float, default=None)
parser.add_argument("--max-samples", type=int, default=None)
parser.add_argument("--log-interval", type=int, default=1)
args = parser.parse_args()
if args.max_samples is not None and args.max_samples <= 0:
raise ValueError("--max-samples must be positive when provided.")
if args.log_interval <= 0:
raise ValueError("--log-interval must be positive.")
config_path = args.config.resolve()
print(
f"inference started: config={config_path} checkpoint={args.checkpoint} device={args.device}",
flush=True,
)
config = load_config(config_path)
validate_config(config)
if args.checkpoint is None:
raise ValueError("Inference requires an explicit --checkpoint path.")
device = resolve_device(args.device)
print(f"runtime: {device_summary(device)}", flush=True)
print("building model", flush=True)
model = build_model(config).to(device)
print(f"model ready: parameters_device={next(model.parameters()).device}", flush=True)
source_root = (
resolve_project_path(args.checkpoint_source_root, config_path)
if args.checkpoint_source_root
else None
)
report = load_checkpoint(
model,
resolve_project_path(args.checkpoint, config_path),
strict=not args.non_strict_checkpoint,
map_location=device,
source_root=source_root,
)
print(f"loaded checkpoint: {report}", flush=True)
data_config = config["data"]
print("building inference dataset", flush=True)
dataset = build_inference_dataset(config, config_path)
loader = DataLoader(
dataset,
batch_size=int(data_config["batch_size"]),
shuffle=False,
num_workers=int(data_config["num_workers"]),
)
print(f"dataset ready: samples={len(dataset)} batches={len(loader)}", flush=True)
output_dir = resolve_project_path(args.output_dir, config_path) if args.output_dir else (
resolve_project_path(config["project"]["output_dir"], config_path) / "inference"
)
mask_ratio = float(config["model"]["mask_ratio"] if args.mask_ratio is None else args.mask_ratio)
if mask_ratio not in {0.0, 0.75}:
raise ValueError("W-MAE supports mask_ratio 0.0 or 0.75 only.")
print(f"inference running: mask_ratio={mask_ratio} output_dir={output_dir}", flush=True)
count = run_inference(model, loader, output_dir, mask_ratio, device, args.max_samples, args.log_interval)
print(f"wrote {count} reconstruction samples to {output_dir}", flush=True)
print("inference completed successfully", flush=True)
if __name__ == "__main__":
main()
|