Upload folder using huggingface_hub

README.md

@@ -0,0 +1,153 @@
+---
+frameworks: PyTorch
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Weather Forecast
+- Masked Autoencoder
+- ERA5
+- W-MAE
+tasks: []
+datasets:
+- OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">W-MAE</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+W-MAE (Weather Masked AutoEncoder) is a pretraining model for multivariable weather forecasting.
+
+Paper: W-MAE: Pre-trained weather model with masked autoencoder for multi-variable weather forecasting
+
+https://arxiv.org/abs/2304.08754
+
+# Model Description
+
+W-MAE first learns spatial relationships among weather variables through masked reconstruction, then learns temporal dependencies by fine-tuning on a forecasting task.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Masked weather-field pretraining | Train the W-MAE reconstruction model with ERA5 HDF5 data that follows this project's protocol. |
+| Local quick validation | Use synthetic HDF5 data to check loading, training, inference, and visualization of inference results. |
+| ModelScope / OneCode execution | Download the standalone model package, configure data, and run the scripts directly. |
+| Multi-GPU training | Launch PyTorch DistributedDataParallel with `torchrun`. |
+
+# Usage Guide
+
+## 1. OneCode Usage
+
+Experience intelligent one-click AI4S programming through the OneCode online environment:
+
+[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- CPU can be used for import and small-scale connectivity verification; full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching your cluster, is recommended.
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/W-MAE --local-dir ./W-MAE
+cd W-MAE
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Please activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+```bash
+# Please activate CONDA first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+### Training Data Introduction
+
+The OneScience community provides an ERA5 data slice for training. Download it to the directory configured by `data.dataset_dir` in `conf/config.yaml`:
+
+```bash
+hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data/era5
+```
+
+The W-MAE adapter expects yearly HDF5 files under `data/era5/data/`. Each file must contain a `fields` dataset, an ordered list of 20 channel names, six-hour time steps, and normalization statistics. Verify the physical ERA5 variable order before scientific training.
+
+### Generate Synthetic Data
+
+When real ERA5 data is unavailable, generate protocol-compatible files for pipeline checks:
+
+```bash
+python scripts/fake_data.py
+```
+
+The synthetic files use placeholder channel names and must not be used for scientific evaluation.
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Multi-GPU:
+
+```bash
+torchrun --nproc_per_node=8 scripts/train.py
+```
+
+Training starts from random initialization and saves `data/checkpoint/model_bak.pth` by default. A compatible checkpoint can be supplied explicitly when continuing training.
+
+### Training Weights
+
+This repository provides a `weight/` directory for W-MAE checkpoints. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+Inference reads `data/checkpoint/model_bak.pth` by default and writes compressed reconstruction samples to `outputs/inference/`:
+
+```bash
+python scripts/inference.py
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+The result script validates reconstruction files and writes diagnostic figures under `outputs/inference/diagnostics/`.
+
+# Official OneScience Resources
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+- This repository is an independent reproduction of the original W-MAE paper.

conf/config.yaml

@@ -0,0 +1,89 @@
+project:
+  name: W-MAE
+  output_dir: ./outputs
+
+data:
+  dataset_dir: ./data/era5
+  train_years: [2000, 2001]
+  val_years: [2002]
+  test_years: [2003]
+  # These names match the generated test HDF5 files. Replace the complete list
+  # with the verified channel order when using real ERA5 data.
+  variables:
+    - synthetic_channel_00
+    - synthetic_channel_01
+    - synthetic_channel_02
+    - synthetic_channel_03
+    - synthetic_channel_04
+    - synthetic_channel_05
+    - synthetic_channel_06
+    - synthetic_channel_07
+    - synthetic_channel_08
+    - synthetic_channel_09
+    - synthetic_channel_10
+    - synthetic_channel_11
+    - synthetic_channel_12
+    - synthetic_channel_13
+    - synthetic_channel_14
+    - synthetic_channel_15
+    - synthetic_channel_16
+    - synthetic_channel_17
+    - synthetic_channel_18
+    - synthetic_channel_19
+  expected_channels: 20
+  input_steps: 1
+  output_steps: 1
+  time_step_hours: 6
+  normalize: true
+  source_size: [721, 1440]
+  model_size: [720, 1440]
+  latitude_crop: drop_last
+  batch_size: 4
+  num_workers: 0
+
+fake_data:
+  years: [2000, 2001, 2002, 2003]
+  time_steps: 40
+  variables:
+    - synthetic_channel_00
+    - synthetic_channel_01
+    - synthetic_channel_02
+    - synthetic_channel_03
+    - synthetic_channel_04
+    - synthetic_channel_05
+    - synthetic_channel_06
+    - synthetic_channel_07
+    - synthetic_channel_08
+    - synthetic_channel_09
+    - synthetic_channel_10
+    - synthetic_channel_11
+    - synthetic_channel_12
+    - synthetic_channel_13
+    - synthetic_channel_14
+    - synthetic_channel_15
+    - synthetic_channel_16
+    - synthetic_channel_17
+    - synthetic_channel_18
+    - synthetic_channel_19
+
+model:
+  image_size: [720, 1440]
+  patch_size: [8, 8]
+  input_channels: 20
+  embed_dim: 768
+  encoder_depth: 12
+  decoder_embed_dim: 512
+  decoder_depth: 6
+  mask_ratio: 0.75
+  norm_pixel_loss: false
+
+inference:
+  output_dir: ./outputs/inference
+  split: test
+  samples: 1
+  mask_ratio: 0.75
+
+validation:
+  expected_channels: 20
+  mask_ratio: 0.75
+  diagnostic_channel: 0

config.json

@@ -0,0 +1,85 @@
+{
+  "model_name": "W-MAE",
+  "model_type": "w_mae",
+  "architectures": [
+    "MaskedAutoencoderAFNO"
+  ],
+  "framework": "PyTorch",
+  "domain": "atmosphere",
+  "task": "masked-weather-pretraining",
+  "implementation": {
+    "entry_point": "model/w_mae.py",
+    "scope": "AFNO-based masked autoencoder for multi-variable weather fields; this metadata describes the local PyTorch implementation and its ERA5 adapter"
+  },
+  "architecture": {
+    "family": "masked autoencoder with Adaptive Fourier Neural Operator (AFNO) blocks",
+    "input_format": "BCHW",
+    "input_grid_shape": [
+      720,
+      1440
+    ],
+    "patch_size": [
+      8,
+      8
+    ],
+    "input_channels": 20,
+    "encoder": {
+      "embed_dim": 768,
+      "depth": 12,
+      "mlp_ratio": 4.0,
+      "afno_num_blocks": 8
+    },
+    "decoder": {
+      "embed_dim": 512,
+      "depth": 6,
+      "mlp_ratio": 4.0
+    },
+    "mask_ratio": 0.75,
+    "supported_mask_ratios": [
+      0.0,
+      0.75
+    ],
+    "normalization": "LayerNorm",
+    "activation": "GELU"
+  },
+  "data": {
+    "dataset": "ERA5",
+    "source_grid_shape": [
+      721,
+      1440
+    ],
+    "preprocessing": "drop the last latitude row to obtain the model grid",
+    "time_step_hours": 6,
+    "input_steps": 1,
+    "output_steps": 1,
+    "channel_names": [
+      "synthetic_channel_00",
+      "synthetic_channel_01",
+      "synthetic_channel_02",
+      "synthetic_channel_03",
+      "synthetic_channel_04",
+      "synthetic_channel_05",
+      "synthetic_channel_06",
+      "synthetic_channel_07",
+      "synthetic_channel_08",
+      "synthetic_channel_09",
+      "synthetic_channel_10",
+      "synthetic_channel_11",
+      "synthetic_channel_12",
+      "synthetic_channel_13",
+      "synthetic_channel_14",
+      "synthetic_channel_15",
+      "synthetic_channel_16",
+      "synthetic_channel_17",
+      "synthetic_channel_18",
+      "synthetic_channel_19"
+    ],
+    "channel_names_note": "The repository configuration intentionally uses synthetic names; physical ERA5 variable names must be supplied and verified before scientific training."
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/w_mae.py",
+    "model/afno.py",
+    "scripts/era5_adapter.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "masked_weather_pretraining",
+  "model": "W-MAE",
+  "input_format": "BCHW",
+  "protocol": "era5_20_channel_720x1440",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/afno.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: int | None = None,
+        out_features: int | None = None,
+        act_layer: type[nn.Module] = nn.GELU,
+        drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        hidden_features = hidden_features or in_features
+        out_features = out_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        return self.drop(x)
+
+
+class AFNO2D(nn.Module):
+    """Adaptive Fourier mixing over a channel-last 2D patch grid."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_blocks: int = 8,
+        sparsity_threshold: float = 0.01,
+        hard_thresholding_fraction: float = 1.0,
+        hidden_size_factor: int = 1,
+    ) -> None:
+        super().__init__()
+        if hidden_size % num_blocks != 0:
+            raise ValueError(f"hidden_size={hidden_size} must be divisible by num_blocks={num_blocks}.")
+        if not 0.0 < hard_thresholding_fraction <= 1.0:
+            raise ValueError("hard_thresholding_fraction must be in (0, 1].")
+
+        self.hidden_size = hidden_size
+        self.sparsity_threshold = sparsity_threshold
+        self.num_blocks = num_blocks
+        self.block_size = hidden_size // num_blocks
+        self.hard_thresholding_fraction = hard_thresholding_fraction
+        self.hidden_size_factor = hidden_size_factor
+        scale = 0.02
+
+        self.w1 = nn.Parameter(
+            scale * torch.randn(2, num_blocks, self.block_size, self.block_size * hidden_size_factor)
+        )
+        self.b1 = nn.Parameter(scale * torch.randn(2, num_blocks, self.block_size * hidden_size_factor))
+        self.w2 = nn.Parameter(
+            scale * torch.randn(2, num_blocks, self.block_size * hidden_size_factor, self.block_size)
+        )
+        self.b2 = nn.Parameter(scale * torch.randn(2, num_blocks, self.block_size))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.ndim != 4 or x.shape[-1] != self.hidden_size:
+            raise ValueError(f"AFNO2D expects [B,H,W,{self.hidden_size}], got {tuple(x.shape)}.")
+
+        bias = x
+        dtype = x.dtype
+        x = torch.fft.rfft2(x.float(), dim=(1, 2), norm="ortho")
+        batch, height, freq_width, _ = x.shape
+        x = x.reshape(batch, height, freq_width, self.num_blocks, self.block_size)
+
+        hidden_block = self.block_size * self.hidden_size_factor
+        o1_real = torch.zeros(
+            batch, height, freq_width, self.num_blocks, hidden_block, device=x.device, dtype=x.real.dtype
+        )
+        o1_imag = torch.zeros_like(o1_real)
+        o2_real = torch.zeros_like(x.real)
+        o2_imag = torch.zeros_like(x.real)
+
+        total_modes = height // 2 + 1
+        kept_modes = max(1, int(total_modes * self.hard_thresholding_fraction))
+        row_slice = slice(total_modes - kept_modes, total_modes + kept_modes)
+        col_slice = slice(0, min(kept_modes, freq_width))
+        selected = x[:, row_slice, col_slice]
+
+        o1_real[:, row_slice, col_slice] = F.relu(
+            torch.einsum("...bi,bio->...bo", selected.real, self.w1[0])
+            - torch.einsum("...bi,bio->...bo", selected.imag, self.w1[1])
+            + self.b1[0]
+        )
+        o1_imag[:, row_slice, col_slice] = F.relu(
+            torch.einsum("...bi,bio->...bo", selected.imag, self.w1[0])
+            + torch.einsum("...bi,bio->...bo", selected.real, self.w1[1])
+            + self.b1[1]
+        )
+
+        hidden_real = o1_real[:, row_slice, col_slice]
+        hidden_imag = o1_imag[:, row_slice, col_slice]
+        o2_real[:, row_slice, col_slice] = (
+            torch.einsum("...bi,bio->...bo", hidden_real, self.w2[0])
+            - torch.einsum("...bi,bio->...bo", hidden_imag, self.w2[1])
+            + self.b2[0]
+        )
+        o2_imag[:, row_slice, col_slice] = (
+            torch.einsum("...bi,bio->...bo", hidden_imag, self.w2[0])
+            + torch.einsum("...bi,bio->...bo", hidden_real, self.w2[1])
+            + self.b2[1]
+        )
+
+        x = torch.stack((o2_real, o2_imag), dim=-1)
+        x = F.softshrink(x, lambd=self.sparsity_threshold)
+        x = torch.view_as_complex(x)
+        x = x.reshape(batch, height, freq_width, self.hidden_size)
+        x = torch.fft.irfft2(x, s=bias.shape[1:3], dim=(1, 2), norm="ortho")
+        return x.to(dtype=dtype) + bias
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        mlp_ratio: float = 4.0,
+        drop: float = 0.0,
+        act_layer: type[nn.Module] = nn.GELU,
+        norm_layer: type[nn.Module] = nn.LayerNorm,
+        double_skip: bool = True,
+        num_blocks: int = 8,
+        sparsity_threshold: float = 0.01,
+        hard_thresholding_fraction: float = 1.0,
+    ) -> None:
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.filter = AFNO2D(dim, num_blocks, sparsity_threshold, hard_thresholding_fraction)
+        self.drop_path = nn.Identity()
+        self.norm2 = norm_layer(dim)
+        self.mlp = Mlp(dim, int(dim * mlp_ratio), act_layer=act_layer, drop=drop)
+        self.double_skip = double_skip
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = self.filter(self.norm1(x))
+        if self.double_skip:
+            x = x + residual
+            residual = x
+        x = self.drop_path(self.mlp(self.norm2(x)))
+        return x + residual

model/check_checkpoint.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+import re
+import subprocess
+import sys
+import zipfile
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+
+REQUIRED_KEYS = {
+    "cls_token",
+    "pos_embed",
+    "mask_token",
+    "decoder_pos_embed",
+    "patch_embed.proj.weight",
+    "decoder_embed.weight",
+    "decoder_norm.weight",
+    "decoder_pred.weight",
+}
+
+PARAMETER_KEY_PATTERN = re.compile(
+    r"(?:decoder_blocks|blocks)\.\d+\.[A-Za-z0-9_.]*?(?:weight|bias|w1|b1|w2|b2)"
+    r"|(?:patch_embed\.proj|decoder_embed|decoder_norm|decoder_pred)\.(?:weight|bias)"
+    r"|(?:cls_token|pos_embed|mask_token|decoder_pos_embed)"
+)
+
+
+def inspect_archive(path: Path) -> dict:
+    if not zipfile.is_zipfile(path):
+        raise ValueError(f"{path} is not a zip-based PyTorch checkpoint.")
+    with zipfile.ZipFile(path) as archive:
+        names = archive.namelist()
+        data_pickle = next((name for name in names if name.endswith("/data.pkl")), None)
+        if data_pickle is None:
+            raise ValueError("Checkpoint archive does not contain data.pkl.")
+        payload = archive.read(data_pickle)
+
+    strings = subprocess.run(
+        ["strings"], input=payload, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True
+    ).stdout.decode("utf-8", errors="ignore").splitlines()
+    parameter_keys = {
+        match.group(0)
+        for line in strings
+        for match in [PARAMETER_KEY_PATTERN.search(line)]
+        if match
+    }
+    encoder_indices = sorted(
+        {int(match.group(1)) for key in parameter_keys if (match := re.match(r"blocks\.(\d+)\.", key))}
+    )
+    decoder_indices = sorted(
+        {int(match.group(1)) for key in parameter_keys if (match := re.match(r"decoder_blocks\.(\d+)\.", key))}
+    )
+    missing = sorted(REQUIRED_KEYS - parameter_keys)
+    return {
+        "path": str(path),
+        "archive_entries": len(names),
+        "parameter_key_count": len(parameter_keys),
+        "encoder_blocks": len(encoder_indices),
+        "decoder_blocks": len(decoder_indices),
+        "missing_required_keys": missing,
+        "status": "compatible_structure" if not missing and len(encoder_indices) == 12 and len(decoder_indices) == 6 else "mismatch",
+    }
+
+
+def inspect_with_torch(path: Path) -> dict:
+    try:
+        import torch
+        from model.w_mae import w_mae_base
+    except (ImportError, OSError) as error:
+        return {"status": "unavailable", "reason": str(error)}
+
+    try:
+        checkpoint = torch.load(path, map_location="cpu")
+    except Exception as error:
+        return {"status": "unavailable", "reason": f"checkpoint deserialization failed: {error}"}
+    state_dict = checkpoint.get("model", checkpoint)
+    model = w_mae_base()
+    model_state = model.state_dict()
+    missing = sorted(set(model_state) - set(state_dict))
+    unexpected = sorted(set(state_dict) - set(model_state))
+    shape_mismatches = {
+        key: {"model": list(model_state[key].shape), "checkpoint": list(state_dict[key].shape)}
+        for key in model_state.keys() & state_dict.keys()
+        if tuple(model_state[key].shape) != tuple(state_dict[key].shape)
+    }
+    return {
+        "status": "compatible" if not missing and not unexpected and not shape_mismatches else "mismatch",
+        "missing_keys": missing,
+        "unexpected_keys": unexpected,
+        "shape_mismatches": shape_mismatches,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Inspect W-MAE checkpoint structure and model compatibility.")
+    parser.add_argument("checkpoint", type=Path)
+    parser.add_argument("--torch-check", action="store_true")
+    args = parser.parse_args()
+
+    report = {"archive": inspect_archive(args.checkpoint)}
+    if args.torch_check:
+        report["torch"] = inspect_with_torch(args.checkpoint)
+    print(json.dumps(report, indent=2))
+    archive_ok = report["archive"]["status"] == "compatible_structure"
+    torch_ok = not args.torch_check or report["torch"]["status"] == "compatible"
+    raise SystemExit(0 if archive_ok and torch_ok else 1)
+
+
+if __name__ == "__main__":
+    main()

model/checkpoint.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from pathlib import Path
+import sys
+from typing import Any, Iterator, Mapping
+
+import torch
+from torch import nn
+
+
+@contextmanager
+def _temporary_module_path(source_root: str | Path | None) -> Iterator[None]:
+    if source_root is None:
+        yield
+        return
+
+    path = str(Path(source_root).resolve())
+    sys.path.insert(0, path)
+    try:
+        yield
+    finally:
+        try:
+            sys.path.remove(path)
+        except ValueError:
+            pass
+
+
+def _extract_state_dict(checkpoint: Any) -> Mapping[str, torch.Tensor]:
+    if isinstance(checkpoint, Mapping):
+        for key in ("model", "model_state", "state_dict"):
+            candidate = checkpoint.get(key)
+            if isinstance(candidate, Mapping):
+                return candidate
+        if checkpoint and all(isinstance(key, str) for key in checkpoint):
+            return checkpoint
+    raise ValueError("Checkpoint does not contain a model state dictionary.")
+
+
+def load_checkpoint(
+    model: nn.Module,
+    checkpoint_path: str | Path,
+    *,
+    strict: bool = True,
+    map_location: str | torch.device = "cpu",
+    source_root: str | Path | None = None,
+) -> dict[str, Any]:
+    """Load a trusted W-MAE checkpoint and return loading metadata.
+
+    ``source_root`` is needed only for legacy checkpoints whose pickle payload
+    references modules from the original repository, such as ``utils.YParams``.
+    """
+    checkpoint_path = Path(checkpoint_path)
+    if not checkpoint_path.is_file():
+        raise FileNotFoundError(f"Checkpoint does not exist: {checkpoint_path}")
+
+    with _temporary_module_path(source_root):
+        checkpoint = torch.load(checkpoint_path, map_location=map_location, weights_only=False)
+
+    state_dict = _extract_state_dict(checkpoint)
+    model_state = model.state_dict()
+    missing = sorted(set(model_state) - set(state_dict))
+    unexpected = sorted(set(state_dict) - set(model_state))
+    shape_mismatches = {
+        key: {"model": list(model_state[key].shape), "checkpoint": list(state_dict[key].shape)}
+        for key in model_state.keys() & state_dict.keys()
+        if tuple(model_state[key].shape) != tuple(state_dict[key].shape)
+    }
+    if strict and (missing or unexpected or shape_mismatches):
+        raise RuntimeError(
+            "Checkpoint is incompatible with the model: "
+            f"missing={missing}, unexpected={unexpected}, shape_mismatches={shape_mismatches}"
+        )
+
+    incompatible = model.load_state_dict(state_dict, strict=strict)
+    return {
+        "checkpoint_path": str(checkpoint_path),
+        "strict": strict,
+        "missing_keys": list(incompatible.missing_keys),
+        "unexpected_keys": list(incompatible.unexpected_keys),
+        "shape_mismatches": shape_mismatches,
+        "checkpoint_keys": len(state_dict),
+    }

model/w_mae.py

@@ -0,0 +1,260 @@
+from __future__ import annotations
+
+from functools import partial
+from typing import NamedTuple
+
+import numpy as np
+import torch
+from torch import nn
+
+from .afno import Block
+
+
+class WMAEOutput(NamedTuple):
+    loss: torch.Tensor
+    prediction: torch.Tensor
+    mask: torch.Tensor
+
+
+def _sincos_1d(embed_dim: int, positions: np.ndarray) -> np.ndarray:
+    if embed_dim % 2 != 0:
+        raise ValueError("The 1D sine-cosine embedding dimension must be even.")
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega = 1.0 / (10000 ** (omega / (embed_dim / 2.0)))
+    values = np.einsum("m,d->md", positions.reshape(-1), omega)
+    return np.concatenate((np.sin(values), np.cos(values)), axis=1)
+
+
+def build_2d_sincos_position_embedding(
+    embed_dim: int, grid_size: tuple[int, int], include_cls_token: bool
+) -> torch.Tensor:
+    if embed_dim % 4 != 0:
+        raise ValueError("The 2D sine-cosine embedding dimension must be divisible by four.")
+    grid_h = np.arange(grid_size[0], dtype=np.float32)
+    grid_w = np.arange(grid_size[1], dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)
+    embedding = np.concatenate(
+        (_sincos_1d(embed_dim // 2, grid[0]), _sincos_1d(embed_dim // 2, grid[1])), axis=1
+    )
+    if include_cls_token:
+        embedding = np.concatenate((np.zeros((1, embed_dim)), embedding), axis=0)
+    return torch.from_numpy(embedding).float().unsqueeze(0)
+
+
+class PatchEmbed(nn.Module):
+    def __init__(
+        self,
+        img_size: tuple[int, int],
+        patch_size: tuple[int, int],
+        in_chans: int,
+        embed_dim: int,
+    ) -> None:
+        super().__init__()
+        if img_size[0] % patch_size[0] or img_size[1] % patch_size[1]:
+            raise ValueError(f"img_size={img_size} must be divisible by patch_size={patch_size}.")
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.ndim != 4 or tuple(x.shape[-2:]) != self.img_size:
+            raise ValueError(f"PatchEmbed expects [B,C,{self.img_size[0]},{self.img_size[1]}], got {tuple(x.shape)}.")
+        return self.proj(x).flatten(2).transpose(1, 2)
+
+
+class MaskedAutoencoderAFNO(nn.Module):
+    """W-MAE pretraining model reconstructed from the official AFNO source."""
+
+    def __init__(
+        self,
+        img_size: tuple[int, int] = (720, 1440),
+        patch_size: tuple[int, int] = (8, 8),
+        in_chans: int = 20,
+        embed_dim: int = 768,
+        depth: int = 12,
+        decoder_embed_dim: int = 512,
+        decoder_depth: int = 6,
+        mlp_ratio: float = 4.0,
+        norm_layer: type[nn.Module] = nn.LayerNorm,
+        norm_pix_loss: bool = False,
+        num_blocks: int = 8,
+        sparsity_threshold: float = 0.01,
+        hard_thresholding_fraction: float = 1.0,
+    ) -> None:
+        super().__init__()
+        self.img_size = tuple(img_size)
+        self.patch_size = tuple(patch_size)
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+        self.decoder_embed_dim = decoder_embed_dim
+        self.norm_pix_loss = norm_pix_loss
+
+        self.patch_embed = PatchEmbed(self.img_size, self.patch_size, in_chans, embed_dim)
+        num_patches = self.patch_embed.num_patches
+        # The official AFNO path keeps this checkpoint key but does not prepend
+        # a class token in forward_encoder, so it must not enter DDP reduction.
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim), requires_grad=False)
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False)
+        block_args = dict(
+            mlp_ratio=mlp_ratio,
+            norm_layer=norm_layer,
+            num_blocks=num_blocks,
+            sparsity_threshold=sparsity_threshold,
+            hard_thresholding_fraction=hard_thresholding_fraction,
+        )
+        self.blocks = nn.ModuleList([Block(dim=embed_dim, **block_args) for _ in range(depth)])
+        self.norm = norm_layer(embed_dim)
+
+        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
+        self.decoder_pos_embed = nn.Parameter(
+            torch.zeros(1, num_patches, decoder_embed_dim), requires_grad=False
+        )
+        decoder_args = dict(block_args)
+        decoder_args["norm_layer"] = norm_layer
+        self.decoder_blocks = nn.ModuleList(
+            [Block(dim=decoder_embed_dim, **decoder_args) for _ in range(decoder_depth)]
+        )
+        self.decoder_norm = norm_layer(decoder_embed_dim)
+        patch_area = self.patch_size[0] * self.patch_size[1]
+        self.decoder_pred = nn.Linear(decoder_embed_dim, in_chans * patch_area)
+        self.initialize_weights()
+
+    def initialize_weights(self) -> None:
+        grid_size = self.patch_embed.grid_size
+        self.pos_embed.data.copy_(build_2d_sincos_position_embedding(self.embed_dim, grid_size, True))
+        self.decoder_pos_embed.data.copy_(
+            build_2d_sincos_position_embedding(self.decoder_embed_dim, grid_size, False)
+        )
+        nn.init.xavier_uniform_(self.patch_embed.proj.weight.data.flatten(1))
+        nn.init.normal_(self.cls_token, std=0.02)
+        nn.init.normal_(self.mask_token, std=0.02)
+        self.apply(self._init_weights)
+
+    @staticmethod
+    def _init_weights(module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.zeros_(module.bias)
+            nn.init.ones_(module.weight)
+
+    def patchify(self, images: torch.Tensor) -> torch.Tensor:
+        if images.ndim != 4 or images.shape[1] != self.in_chans or tuple(images.shape[-2:]) != self.img_size:
+            raise ValueError(
+                f"patchify expects [B,{self.in_chans},{self.img_size[0]},{self.img_size[1]}], "
+                f"got {tuple(images.shape)}."
+            )
+        ph, pw = self.patch_size
+        gh, gw = self.patch_embed.grid_size
+        x = images.reshape(images.shape[0], self.in_chans, gh, ph, gw, pw)
+        x = torch.einsum("nchpwq->nhwpqc", x)
+        return x.reshape(images.shape[0], gh * gw, ph * pw * self.in_chans)
+
+    def unpatchify(self, patches: torch.Tensor) -> torch.Tensor:
+        ph, pw = self.patch_size
+        gh, gw = self.patch_embed.grid_size
+        expected_dim = ph * pw * self.in_chans
+        if patches.ndim != 3 or patches.shape[1:] != (gh * gw, expected_dim):
+            raise ValueError(f"unpatchify expects [B,{gh * gw},{expected_dim}], got {tuple(patches.shape)}.")
+        x = patches.reshape(patches.shape[0], gh, gw, ph, pw, self.in_chans)
+        x = torch.einsum("nhwpqc->nchpwq", x)
+        return x.reshape(patches.shape[0], self.in_chans, gh * ph, gw * pw)
+
+    @staticmethod
+    def random_masking(
+        x: torch.Tensor, mask_ratio: float
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if mask_ratio not in {0.0, 0.75}:
+            raise ValueError("Official W-MAE AFNO grid reshaping supports mask_ratio 0.0 or 0.75 only.")
+        batch, length, channels = x.shape
+        len_keep = int(length * (1.0 - mask_ratio))
+        noise = torch.rand(batch, length, device=x.device)
+        ids_shuffle = torch.argsort(noise, dim=1)
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, 1, ids_keep.unsqueeze(-1).expand(-1, -1, channels))
+        mask = torch.ones(batch, length, device=x.device)
+        mask[:, :len_keep] = 0
+        return x_masked, torch.gather(mask, 1, ids_restore), ids_restore
+
+    def forward_encoder(
+        self, images: torch.Tensor, mask_ratio: float
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        x = self.patch_embed(images) + self.pos_embed[:, 1:]
+        x, mask, ids_restore = self.random_masking(x, mask_ratio)
+        grid_h, grid_w = self.patch_embed.grid_size
+        divisor = 2 if mask_ratio == 0.75 else 1
+        expected_tokens = (grid_h // divisor) * (grid_w // divisor)
+        if x.shape[1] != expected_tokens:
+            raise ValueError("The selected mask ratio does not form the rectangular AFNO grid expected by W-MAE.")
+        x = x.reshape(x.shape[0], grid_h // divisor, grid_w // divisor, self.embed_dim)
+        for block in self.blocks:
+            x = block(x)
+        x = self.norm(x)
+        return x.flatten(1, 2), mask, ids_restore
+
+    def forward_decoder(self, latent: torch.Tensor, ids_restore: torch.Tensor) -> torch.Tensor:
+        x = self.decoder_embed(latent)
+        missing_tokens = ids_restore.shape[1] - x.shape[1]
+        if missing_tokens < 0:
+            raise ValueError("Latent token count exceeds the decoder target token count.")
+        mask_tokens = self.mask_token.expand(x.shape[0], missing_tokens, -1)
+        x = torch.cat((x, mask_tokens), dim=1)
+        x = torch.gather(x, 1, ids_restore.unsqueeze(-1).expand(-1, -1, x.shape[-1]))
+        x = x + self.decoder_pos_embed
+        grid_h, grid_w = self.patch_embed.grid_size
+        x = x.reshape(x.shape[0], grid_h, grid_w, self.decoder_embed_dim)
+        for block in self.decoder_blocks:
+            x = block(x)
+        x = self.decoder_pred(self.decoder_norm(x))
+        return x.flatten(1, 2)
+
+    def forward_loss(
+        self, images: torch.Tensor, prediction: torch.Tensor, mask: torch.Tensor, mask_ratio: float
+    ) -> torch.Tensor:
+        target = self.patchify(images)
+        if self.norm_pix_loss:
+            mean = target.mean(dim=-1, keepdim=True)
+            variance = target.var(dim=-1, keepdim=True)
+            target = (target - mean) / torch.sqrt(variance + 1e-6)
+        loss = (prediction - target).pow(2).mean(dim=-1)
+        if mask_ratio == 0.0:
+            return loss.mean()
+        masked_count = mask.sum()
+        if masked_count.item() == 0:
+            raise ValueError("Masked reconstruction loss requires at least one masked patch.")
+        return (loss * mask).sum() / masked_count
+
+    def forward(self, images: torch.Tensor, mask_ratio: float = 0.75) -> WMAEOutput:
+        latent, mask, ids_restore = self.forward_encoder(images, mask_ratio)
+        prediction = self.forward_decoder(latent, ids_restore)
+        loss = self.forward_loss(images, prediction, mask, mask_ratio)
+        return WMAEOutput(loss, prediction, mask)
+
+
+def w_mae_base(
+    embed_dim: int = 768,
+    depth: int = 12,
+    decoder_embed_dim: int = 512,
+    decoder_depth: int = 6,
+    mlp_ratio: float = 4.0,
+    norm_layer: type[nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+    **kwargs: object,
+) -> MaskedAutoencoderAFNO:
+    return MaskedAutoencoderAFNO(
+        embed_dim=embed_dim,
+        depth=depth,
+        decoder_embed_dim=decoder_embed_dim,
+        decoder_depth=decoder_depth,
+        mlp_ratio=mlp_ratio,
+        norm_layer=norm_layer,
+        **kwargs,
+    )
+
+
+mae_vit_base_patch16 = w_mae_base

scripts/era5_adapter.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Sequence
+
+import h5py
+
+
+EXPECTED_CHANNELS = 20
+SOURCE_SIZE = (721, 1440)
+MODEL_SIZE = (720, 1440)
+
+
+def _decode_variables(values: Sequence[Any]) -> list[str]:
+    return [value.decode() if isinstance(value, bytes) else str(value) for value in values]
+
+
+def validate_variables(variables: Sequence[str], expected_channels: int = EXPECTED_CHANNELS) -> list[str]:
+    names = list(variables)
+    if len(names) != expected_channels:
+        raise ValueError(
+            f"W-MAE requires exactly {expected_channels} explicitly ordered variables; "
+            f"received {len(names)}. The official repository does not publish a complete "
+            "channel-name mapping, so no default mapping is assumed."
+        )
+    if len(set(names)) != len(names):
+        raise ValueError("W-MAE variable names must be unique and explicitly ordered.")
+    return names
+
+
+def inspect_era5_contract(
+    dataset_dir: str | Path,
+    years: Sequence[int],
+    variables: Sequence[str],
+    expected_channels: int = EXPECTED_CHANNELS,
+) -> dict[str, Any]:
+    """Validate the HDF5 metadata before importing the torch-based datapipe."""
+    dataset_dir = Path(dataset_dir)
+    names = validate_variables(variables, expected_channels)
+    missing_years = [year for year in years if not (dataset_dir / "data" / f"{year}.h5").is_file()]
+    if missing_years:
+        raise ValueError(f"ERA5 year files are missing: {missing_years}")
+
+    first_file = dataset_dir / "data" / f"{years[0]}.h5"
+    with h5py.File(first_file, "r") as handle:
+        if "fields" not in handle:
+            raise ValueError(f"{first_file} does not contain a 'fields' dataset.")
+        fields = handle["fields"]
+        if fields.ndim != 4:
+            raise ValueError(f"fields must have shape [T,C,H,W], got {fields.shape}.")
+        if tuple(fields.shape[-2:]) != SOURCE_SIZE:
+            raise ValueError(f"W-MAE expects source grid {SOURCE_SIZE}, got {fields.shape[-2:]}.")
+        if "variables" not in fields.attrs or "time_step" not in fields.attrs:
+            raise ValueError("fields must define 'variables' and 'time_step' attributes.")
+        available = _decode_variables(fields.attrs["variables"])
+        fields_shape = tuple(fields.shape)
+        time_step = int(fields.attrs["time_step"])
+        missing_variables = [name for name in names if name not in available]
+        if missing_variables:
+            raise ValueError(f"Configured variables are absent from HDF5 metadata: {missing_variables}")
+        if "global_means" not in handle or "global_stds" not in handle:
+            stats_dir = dataset_dir / "stats"
+            if not (stats_dir / "global_means.npy").is_file() or not (stats_dir / "global_stds.npy").is_file():
+                raise ValueError("Normalization statistics are missing from HDF5 and dataset_dir/stats.")
+
+    return {
+        "file": str(first_file),
+        "fields_shape": fields_shape,
+        "time_step": time_step,
+        "selected_variables": names,
+        "channel_indices": [available.index(name) for name in names],
+        "crop": "fields[..., :720, :]",
+        "model_size": MODEL_SIZE,
+    }
+
+
+def _crop_last_latitude(value: Any) -> Any:
+    if tuple(value.shape[-2:]) == MODEL_SIZE:
+        return value
+    if tuple(value.shape[-2:]) != SOURCE_SIZE:
+        raise ValueError(f"Expected trailing spatial shape {SOURCE_SIZE}, got {tuple(value.shape[-2:])}.")
+    return value[..., : MODEL_SIZE[0], :]
+
+
+class WMAEERA5Dataset:
+    """Thin W-MAE adapter around OneScience ERA5Dataset.
+
+    The official W-MAE loader removes the final latitude row from a
+    721x1440 ERA5 field. This wrapper preserves OneScience loading and
+    normalization while applying that exact spatial convention.
+    """
+
+    def __init__(
+        self,
+        dataset_dir: str | Path,
+        years: Sequence[int],
+        variables: Sequence[str],
+        task: str = "forecast",
+        input_steps: int = 1,
+        output_steps: int = 1,
+        normalize: bool = True,
+    ) -> None:
+        if task not in {"pretrain", "forecast"}:
+            raise ValueError("task must be either 'pretrain' or 'forecast'.")
+        names = validate_variables(variables)
+        inspect_era5_contract(dataset_dir, years, names)
+
+        try:
+            from onescience.datapipes.climate.era5 import ERA5Dataset
+        except (ImportError, OSError) as error:
+            raise RuntimeError(
+                "OneScience ERA5Dataset could not be imported. Verify the active "
+                "OneScience/PyTorch runtime before constructing WMAEERA5Dataset."
+            ) from error
+
+        self.task = task
+        self.dataset = ERA5Dataset(
+            dataset_dir=str(dataset_dir),
+            used_years=list(years),
+            used_variables=names,
+            input_steps=input_steps,
+            output_steps=output_steps,
+            normalize=normalize,
+        )
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def __getitem__(self, index: int) -> tuple[Any, Any, Any, int, list[str]]:
+        invar, outvar, cos_zenith, step_idx, time_index = self.dataset[index]
+        invar = _crop_last_latitude(invar)
+        outvar = _crop_last_latitude(outvar)
+        cos_zenith = _crop_last_latitude(cos_zenith)
+        target = invar if self.task == "pretrain" else outvar
+        return invar, target, cos_zenith, step_idx, time_index

scripts/fake_data.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import sys
+
+import h5py
+import numpy as np
+import yaml
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+sys.path.insert(0, str(PROJECT_ROOT / "scripts"))
+
+from era5_adapter import inspect_era5_contract
+
+
+def load_config(path: Path) -> dict:
+    with path.open("r", encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def generate_year_file(
+    output_path: Path,
+    variables: list[str],
+    time_steps: int,
+    height: int,
+    width: int,
+    time_step_hours: int,
+) -> None:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    channels = len(variables)
+    means = np.zeros((1, channels, 1, 1), dtype=np.float32)
+    stds = np.ones((1, channels, 1, 1), dtype=np.float32)
+
+    with h5py.File(output_path, "w") as handle:
+        fields = handle.create_dataset(
+            "fields",
+            shape=(time_steps, channels, height, width),
+            dtype="float32",
+            chunks=(1, 1, height, width),
+            fillvalue=0.0,
+        )
+        fields.attrs["variables"] = variables
+        fields.attrs["time_step"] = time_step_hours
+        handle.create_dataset("global_means", data=means)
+        handle.create_dataset("global_stds", data=stds)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Generate lightweight ERA5-style HDF5 files for W-MAE.")
+    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "conf" / "config.yaml")
+    parser.add_argument("--output-dir", type=Path, default=None)
+    args = parser.parse_args()
+
+    config = load_config(args.config)
+    data_config = config["data"]
+    fake_config = config["fake_data"]
+    output_dir = args.output_dir or PROJECT_ROOT / data_config["dataset_dir"]
+    variables = list(fake_config["variables"])
+    years = list(fake_config["years"])
+    source_height, source_width = data_config["source_size"]
+
+    for year in years:
+        path = output_dir / "data" / f"{year}.h5"
+        generate_year_file(
+            output_path=path,
+            variables=variables,
+            time_steps=int(fake_config["time_steps"]),
+            height=int(source_height),
+            width=int(source_width),
+            time_step_hours=int(data_config["time_step_hours"]),
+        )
+        print(f"created {path} ({path.stat().st_size / 1024:.1f} KiB physical size)")
+
+    report = inspect_era5_contract(output_dir, years, variables)
+    print(f"validated fields shape: {report['fields_shape']}")
+    print(f"W-MAE spatial adapter: {report['crop']} -> {report['model_size']}")
+    print("Synthetic channel names are test-only and must not be used for real training.")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,177 @@
+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()

scripts/result.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_RESULTS_DIR = PROJECT_ROOT / "outputs" / "inference"
+
+
+def _load_array(path: Path, name: str) -> np.ndarray:
+    value = np.load(path)[name]
+    if not np.all(np.isfinite(value)):
+        raise ValueError(f"{path.name}:{name} contains NaN or Inf values.")
+    return value
+
+
+def validate_sample(path: Path, expected_channels: int = 20) -> dict[str, Any]:
+    with np.load(path) as result:
+        required = {"input", "reconstruction", "error", "mask", "step_idx", "time_index"}
+        missing = required - set(result.files)
+        if missing:
+            raise ValueError(f"{path} is missing fields: {sorted(missing)}")
+        inputs = _load_array(path, "input")
+        reconstruction = _load_array(path, "reconstruction")
+        error = _load_array(path, "error")
+        mask = _load_array(path, "mask")
+
+    expected_shape = (expected_channels, 720, 1440)
+    if inputs.shape != expected_shape:
+        raise ValueError(f"input shape must be {expected_shape}, got {inputs.shape}.")
+    if reconstruction.shape != inputs.shape or error.shape != inputs.shape:
+        raise ValueError("reconstruction and error must have the same shape as input.")
+    if mask.shape != (90, 180):
+        raise ValueError(f"mask must be a patch grid with shape (90, 180), got {mask.shape}.")
+    if not np.array_equal(mask, mask.astype(bool)):
+        raise ValueError("mask must contain only binary values.")
+    if not np.allclose(error, reconstruction - inputs):
+        raise ValueError("error does not equal reconstruction - input.")
+
+    spatial_mask = np.repeat(np.repeat(mask.astype(bool), 8, axis=0), 8, axis=1)
+    visible_error = np.abs(error[:, ~spatial_mask])
+    masked_error = np.abs(error[:, spatial_mask])
+    mask_fraction = float(mask.mean())
+    return {
+        "file": str(path),
+        "shape": list(inputs.shape),
+        "mask_fraction": mask_fraction,
+        "masked_pixels": int(spatial_mask.sum()),
+        "visible_mae": np.mean(visible_error, axis=1).tolist() if visible_error.size else [0.0] * expected_channels,
+        "masked_mae": np.mean(masked_error, axis=1).tolist() if masked_error.size else [0.0] * expected_channels,
+        "mae": np.mean(np.abs(error), axis=(1, 2)).tolist(),
+        "rmse": np.sqrt(np.mean(error**2, axis=(1, 2))).tolist(),
+    }
+
+
+def save_diagnostic_png(path: Path, output_path: Path, channel: int) -> None:
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    with np.load(path) as result:
+        input_field = result["input"][channel]
+        reconstruction = result["reconstruction"][channel]
+        error = result["error"][channel]
+        mask = result["mask"].astype(bool)
+        time_index = result["time_index"].item()
+
+    input_low, input_high = np.nanpercentile(input_field, [1, 99])
+    reconstruction_low, reconstruction_high = np.nanpercentile(reconstruction, [1, 99])
+    field_low = min(input_low, reconstruction_low)
+    field_high = max(input_high, reconstruction_high)
+    if not np.isfinite(field_low) or not np.isfinite(field_high) or field_high <= field_low:
+        field_low, field_high = float(np.nanmin(input_field)), float(np.nanmax(input_field))
+    error_limit = float(np.nanpercentile(np.abs(error), 99))
+    if not np.isfinite(error_limit) or error_limit <= 0:
+        error_limit = max(float(np.nanmax(np.abs(error))), 1.0)
+
+    latitude = np.linspace(90, -90, input_field.shape[0])
+    longitude = np.linspace(0, 360, input_field.shape[1], endpoint=False)
+    extent = [longitude[0], longitude[-1], latitude[-1], latitude[0]]
+    masked_pixels = np.repeat(np.repeat(mask, 8, axis=0), 8, axis=1)
+    masked_overlay = np.ma.masked_where(~masked_pixels, masked_pixels)
+
+    figure, axes = plt.subplots(2, 3, figsize=(18, 8), constrained_layout=True)
+    field_kwargs = {"cmap": "viridis", "vmin": field_low, "vmax": field_high, "extent": extent, "aspect": "auto"}
+    error_kwargs = {"cmap": "RdBu_r", "vmin": -error_limit, "vmax": error_limit, "extent": extent, "aspect": "auto"}
+    panels = [
+        (axes[0, 0], input_field, "Input (normalized)", field_kwargs),
+        (axes[0, 1], reconstruction, "Reconstruction (normalized)", field_kwargs),
+        (axes[0, 2], error, "Signed error", error_kwargs),
+        (axes[1, 0], np.abs(error), "Absolute error", {**error_kwargs, "cmap": "magma", "vmin": 0}),
+    ]
+    for axis, field, title, kwargs in panels:
+        image = axis.imshow(field, **kwargs)
+        axis.set_title(title)
+        axis.set_xlabel("Longitude (degrees east)")
+        axis.set_ylabel("Latitude (degrees)")
+        axis.contour(masked_pixels, levels=[0.5], colors="white", linewidths=0.35, extent=extent)
+        figure.colorbar(image, ax=axis, shrink=0.8)
+
+    mask_image = axes[1, 1].imshow(mask.astype(float), cmap="gray_r", vmin=0, vmax=1, extent=extent, aspect="auto")
+    axes[1, 1].set_title("Patch mask (white = masked)")
+    axes[1, 1].set_xlabel("Longitude (degrees east)")
+    axes[1, 1].set_ylabel("Latitude (degrees)")
+    figure.colorbar(mask_image, ax=axes[1, 1], ticks=[0, 1], shrink=0.8)
+
+    axes[1, 2].hist(error.ravel(), bins=80, color="#315f8c", alpha=0.85)
+    axes[1, 2].axvline(0, color="black", linewidth=0.8)
+    axes[1, 2].set_title(f"Error distribution\nMAE={np.mean(np.abs(error)):.4f}, RMSE={np.sqrt(np.mean(error**2)):.4f}")
+    axes[1, 2].set_xlabel("Signed error (normalized)")
+    axes[1, 2].set_ylabel("Pixel count")
+
+    figure.suptitle(f"W-MAE reconstruction diagnostic | channel={channel} | time={time_index}")
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    figure.savefig(output_path, dpi=150)
+    plt.close(figure)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Validate W-MAE inference outputs and make diagnostics.")
+    parser.add_argument(
+        "results_dir",
+        type=Path,
+        nargs="?",
+        default=DEFAULT_RESULTS_DIR,
+        help="Inference output directory (default: outputs/inference).",
+    )
+    parser.add_argument("--expected-channels", type=int, default=20)
+    parser.add_argument("--mask-ratio", type=float, default=None)
+    parser.add_argument(
+        "--channel",
+        type=int,
+        default=0,
+        help="Channel for the diagnostic image (default: 0). Use --channel -1 to disable it.",
+    )
+    parser.add_argument("--diagnostic-dir", type=Path, default=None)
+    parser.add_argument(
+        "--visualization-samples",
+        type=int,
+        default=1,
+        help="Number of samples to visualize; validation still covers all samples (default: 1).",
+    )
+    args = parser.parse_args()
+    if args.visualization_samples <= 0:
+        raise ValueError("--visualization-samples must be positive.")
+    files = sorted(args.results_dir.glob("sample_*.npz"))
+    if not files:
+        raise FileNotFoundError(f"No sample_*.npz files found in {args.results_dir}.")
+    reports = [validate_sample(path, args.expected_channels) for path in files]
+    if args.mask_ratio is not None:
+        for report in reports:
+            if not np.isclose(report["mask_fraction"], args.mask_ratio):
+                raise ValueError(
+                    f"{report['file']} mask fraction {report['mask_fraction']} "
+                    f"does not match expected {args.mask_ratio}."
+                )
+    if args.channel is not None and args.channel >= 0:
+        if not 0 <= args.channel < args.expected_channels:
+            raise ValueError("--channel is outside the configured channel range.")
+        diagnostic_dir = args.diagnostic_dir or args.results_dir / "diagnostics"
+        for path in files[:args.visualization_samples]:
+            save_diagnostic_png(path, diagnostic_dir / f"{path.stem}_channel_{args.channel:02d}.png", args.channel)
+    summary = {
+        "samples": len(reports),
+        "mean_mask_fraction": float(np.mean([report["mask_fraction"] for report in reports])),
+        "mean_mae": np.mean([report["mae"] for report in reports], axis=0).tolist(),
+        "mean_rmse": np.mean([report["rmse"] for report in reports], axis=0).tolist(),
+        "reports": reports,
+    }
+    print(json.dumps(summary, indent=2), flush=True)
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,384 @@
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+import sys
+import traceback
+from typing import Any
+
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+import yaml
+
+
+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 model.w_mae import w_mae_base
+from era5_adapter import WMAEERA5Dataset
+
+
+def load_config(path: Path) -> dict[str, Any]:
+    with path.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    if not isinstance(config, dict):
+        raise ValueError(f"Configuration must contain a mapping: {path}")
+    return config
+
+
+def resolve_project_path(value: str | Path, config_path: Path) -> Path:
+    path = Path(value)
+    return path if path.is_absolute() else (config_path.parent.parent / path).resolve()
+
+
+def validate_config(config: dict[str, Any]) -> None:
+    data = config["data"]
+    model = config["model"]
+    if int(data["expected_channels"]) != int(model["input_channels"]):
+        raise ValueError("data.expected_channels must equal model.input_channels.")
+    if tuple(data["model_size"]) != tuple(model["image_size"]):
+        raise ValueError("data.model_size must equal model.image_size.")
+    if tuple(model["patch_size"])[0] <= 0 or tuple(model["patch_size"])[1] <= 0:
+        raise ValueError("model.patch_size values must be positive.")
+    if float(model["mask_ratio"]) not in {0.0, 0.75}:
+        raise ValueError("W-MAE supports mask_ratio 0.0 or 0.75 only.")
+
+
+def build_model(config: dict[str, Any]) -> torch.nn.Module:
+    model_config = config["model"]
+    return w_mae_base(
+        img_size=tuple(model_config["image_size"]),
+        patch_size=tuple(model_config["patch_size"]),
+        in_chans=int(model_config["input_channels"]),
+        embed_dim=int(model_config["embed_dim"]),
+        depth=int(model_config["encoder_depth"]),
+        decoder_embed_dim=int(model_config["decoder_embed_dim"]),
+        decoder_depth=int(model_config["decoder_depth"]),
+        mlp_ratio=float(model_config.get("mlp_ratio", 4.0)),
+        norm_pix_loss=bool(model_config["norm_pixel_loss"]),
+    )
+
+
+def build_dataset(config: dict[str, Any], config_path: Path, split: str) -> WMAEERA5Dataset:
+    data = config["data"]
+    years = data[f"{split}_years"]
+    if not data["variables"]:
+        raise ValueError(
+            "data.variables is empty. Supply the verified, explicitly ordered 20-channel list before training."
+        )
+    return WMAEERA5Dataset(
+        dataset_dir=resolve_project_path(data["dataset_dir"], config_path),
+        years=years,
+        variables=data["variables"],
+        task="pretrain",
+        input_steps=int(data["input_steps"]),
+        output_steps=int(data["output_steps"]),
+        normalize=bool(data["normalize"]),
+    )
+
+
+def validate_architecture_reference(
+    model: torch.nn.Module,
+    reference_path: Path,
+    config_path: Path,
+    source_root: Path | None,
+) -> None:
+    reference_model = build_model(load_config(config_path))
+    load_checkpoint(reference_model, reference_path, source_root=source_root)
+    model_state = model.state_dict()
+    reference_state = reference_model.state_dict()
+    if set(model_state) != set(reference_state):
+        raise RuntimeError("Model parameter names do not match the official architecture reference.")
+    mismatches = {
+        key: (tuple(model_state[key].shape), tuple(reference_state[key].shape))
+        for key in model_state
+        if tuple(model_state[key].shape) != tuple(reference_state[key].shape)
+    }
+    if mismatches:
+        raise RuntimeError(f"Model parameter shapes do not match the official architecture: {mismatches}")
+
+
+def save_checkpoint(
+    path: Path,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    epoch: int,
+    loss: float,
+    config: dict[str, Any],
+) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(
+        {
+            "model": model.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "epoch": epoch,
+            "loss": loss,
+            "config": config,
+        },
+        path,
+    )
+
+
+def resolve_device(requested: str) -> torch.device:
+    if requested == "auto":
+        requested = "cuda" if torch.cuda.is_available() else "cpu"
+    device = torch.device(requested)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("The requested cuda device is unavailable. Use --device cpu or check the DCU runtime.")
+    return device
+
+
+def device_summary(device: torch.device) -> str:
+    if device.type != "cuda":
+        return f"device={device}"
+    backend = f"HIP {torch.version.hip}" if torch.version.hip else f"CUDA {torch.version.cuda}"
+    return f"device={device} name={torch.cuda.get_device_name(device)} backend={backend}"
+
+
+def memory_summary(device: torch.device) -> str:
+    if device.type != "cuda":
+        return "memory=unavailable(cpu)"
+    allocated = torch.cuda.memory_allocated(device) / (1024**3)
+    reserved = torch.cuda.memory_reserved(device) / (1024**3)
+    peak = torch.cuda.max_memory_allocated(device) / (1024**3)
+    return f"memory_allocated={allocated:.2f}GB memory_reserved={reserved:.2f}GB peak_allocated={peak:.2f}GB"
+
+
+def setup_distributed(requested_device: str) -> tuple[torch.device, int, int, int]:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    rank = int(os.environ.get("RANK", "0"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    if world_size == 1:
+        return resolve_device(requested_device), rank, world_size, local_rank
+
+    if requested_device == "cpu":
+        device = torch.device("cpu")
+        backend = "gloo"
+    else:
+        if not torch.cuda.is_available():
+            raise RuntimeError("torchrun requires one accessible DCU/GPU per local rank.")
+        if local_rank >= torch.cuda.device_count():
+            raise RuntimeError(
+                f"LOCAL_RANK={local_rank} exceeds the {torch.cuda.device_count()} visible accelerator devices."
+            )
+        torch.cuda.set_device(local_rank)
+        device = torch.device("cuda", local_rank)
+        backend = "nccl"
+    dist.init_process_group(backend=backend, init_method="env://")
+    return device, rank, world_size, local_rank
+
+
+def rank_print(message: str, rank: int) -> None:
+    if rank == 0:
+        print(message, flush=True)
+
+
+def all_ranks_finite(value: torch.Tensor, world_size: int) -> bool:
+    finite = torch.isfinite(value).to(dtype=torch.int32)
+    if world_size > 1:
+        dist.all_reduce(finite, op=dist.ReduceOp.MIN)
+    return bool(finite.item())
+
+
+def unwrap_model(model: torch.nn.Module) -> torch.nn.Module:
+    return model.module if isinstance(model, DistributedDataParallel) else model
+
+
+def train_one_epoch(
+    model: torch.nn.Module,
+    loader: DataLoader,
+    optimizer: torch.optim.Optimizer,
+    mask_ratio: float,
+    device: torch.device,
+    epoch: int,
+    log_interval: int,
+    rank: int,
+    world_size: int,
+) -> float:
+    model.train()
+    total_loss = 0.0
+    batches = 0
+    total_batches = len(loader)
+    if device.type == "cuda":
+        torch.cuda.reset_peak_memory_stats(device)
+    for batch_index, (inputs, _, _, _, _) in enumerate(loader, start=1):
+        inputs = inputs.to(device)
+        output = model(inputs, mask_ratio=mask_ratio)
+        if not all_ranks_finite(output.loss.detach(), world_size):
+            raise FloatingPointError("Training loss is NaN or Inf on at least one rank.")
+        optimizer.zero_grad(set_to_none=True)
+        output.loss.backward()
+        optimizer.step()
+        total_loss += float(output.loss.detach().cpu())
+        batches += 1
+        if batch_index == 1 or batch_index % log_interval == 0 or batch_index == total_batches:
+            display_loss = output.loss.detach()
+            if world_size > 1:
+                dist.all_reduce(display_loss, op=dist.ReduceOp.SUM)
+                display_loss /= world_size
+            rank_print(
+                f"epoch={epoch} batch={batch_index}/{total_batches} loss={float(display_loss):.6f} "
+                f"{memory_summary(device)}",
+                rank,
+            )
+    if batches == 0:
+        raise ValueError("Training DataLoader produced no batches.")
+    totals = torch.tensor([total_loss, batches], dtype=torch.float64, device=device)
+    if world_size > 1:
+        dist.all_reduce(totals, op=dist.ReduceOp.SUM)
+    return float((totals[0] / totals[1]).cpu())
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train W-MAE using the configured ERA5 adapter.")
+    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "conf" / "config.yaml")
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--learning-rate", type=float, default=1e-4)
+    parser.add_argument(
+        "--device",
+        default="auto",
+        help="Training device: auto (default), cuda (AMD DCU through HIP), or cpu.",
+    )
+    parser.add_argument("--checkpoint", type=Path, default=None)
+    parser.add_argument("--checkpoint-source-root", type=Path, default=None)
+    parser.add_argument("--non-strict-checkpoint", action="store_true")
+    parser.add_argument(
+        "--output-checkpoint",
+        type=Path,
+        default=PROJECT_ROOT / "data" / "checkpoint" / "model_bak.pth",
+    )
+    parser.add_argument("--architecture-reference-checkpoint", type=Path, default=None)
+    parser.add_argument("--log-interval", type=int, default=1)
+    args = parser.parse_args()
+    if args.epochs <= 0:
+        raise ValueError("--epochs must be positive.")
+    if args.log_interval <= 0:
+        raise ValueError("--log-interval must be positive.")
+
+    config_path = args.config.resolve()
+    device, rank, world_size, local_rank = setup_distributed(args.device)
+    rank_print(
+        f"training started: config={config_path} epochs={args.epochs} device={args.device}", rank
+    )
+    rank_print(
+        f"distributed: enabled={world_size > 1} world_size={world_size} "
+        f"rank={rank} local_rank={local_rank}",
+        rank,
+    )
+    config = load_config(config_path)
+    validate_config(config)
+    data_config = config["data"]
+    print(f"rank={rank} local_rank={local_rank} runtime: {device_summary(device)}", flush=True)
+    rank_print("building model", rank)
+    model = build_model(config).to(device)
+    print(
+        f"rank={rank} local_rank={local_rank} model_parameters_device="
+        f"{next(model.parameters()).device}",
+        flush=True,
+    )
+
+    if args.checkpoint is not None:
+        checkpoint_path = resolve_project_path(args.checkpoint, config_path)
+        source_root = (
+            resolve_project_path(args.checkpoint_source_root, config_path)
+            if args.checkpoint_source_root
+            else None
+        )
+        report = load_checkpoint(
+            model,
+            checkpoint_path,
+            strict=not args.non_strict_checkpoint,
+            map_location=device,
+            source_root=source_root,
+        )
+        rank_print(f"loaded checkpoint: {report}", rank)
+    else:
+        rank_print("training from randomly initialized weights", rank)
+
+    if world_size > 1:
+        if device.type == "cuda":
+            model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
+        else:
+            model = DistributedDataParallel(model)
+    rank_print("building dataset", rank)
+    dataset = build_dataset(config, config_path, "train")
+    sampler = DistributedSampler(dataset, shuffle=True) if world_size > 1 else None
+    loader = DataLoader(
+        dataset,
+        batch_size=int(data_config["batch_size"]),
+        shuffle=sampler is None,
+        sampler=sampler,
+        num_workers=int(data_config["num_workers"]),
+    )
+    rank_print(
+        f"dataset ready: samples={len(dataset)} batches_per_rank={len(loader)} "
+        f"batch_size_per_rank={data_config['batch_size']} "
+        f"global_batch_size={int(data_config['batch_size']) * world_size}",
+        rank,
+    )
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate)
+    for epoch in range(args.epochs):
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        rank_print(f"epoch={epoch + 1}/{args.epochs} started", rank)
+        loss = train_one_epoch(
+            model,
+            loader,
+            optimizer,
+            float(config["model"]["mask_ratio"]),
+            device,
+            epoch + 1,
+            args.log_interval,
+            rank,
+            world_size,
+        )
+        rank_print(f"epoch={epoch + 1}/{args.epochs} finished loss={loss:.6f}", rank)
+
+    if world_size > 1:
+        dist.barrier()
+
+    if rank != 0:
+        return
+
+    checkpoint_model = unwrap_model(model)
+
+    reference_path = (
+        resolve_project_path(args.architecture_reference_checkpoint, config_path)
+        if args.architecture_reference_checkpoint
+        else None
+    )
+    reference_source_root = (
+        resolve_project_path(args.checkpoint_source_root, config_path)
+        if args.checkpoint_source_root
+        else None
+    )
+    if reference_path is not None:
+        print(f"checking architecture against {reference_path}", flush=True)
+        validate_architecture_reference(checkpoint_model, reference_path, config_path, reference_source_root)
+        print("architecture reference check passed", flush=True)
+    else:
+        print("architecture reference check skipped: no reference checkpoint supplied", flush=True)
+    output_path = resolve_project_path(args.output_checkpoint, config_path)
+    save_checkpoint(output_path, checkpoint_model, optimizer, args.epochs, loss, config)
+    print(f"checkpoint saved: {output_path}", flush=True)
+    print("training completed successfully", flush=True)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception:
+        traceback.print_exc()
+        raise
+    finally:
+        if dist.is_available() and dist.is_initialized():
+            dist.destroy_process_group()

weight/.gitkeep

@@ -0,0 +1 @@
+# Reserved for model weights.