Upload folder using huggingface_hub

README.md

@@ -0,0 +1,137 @@
+---
+frameworks: PyTorch
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Ocean Forecasting
+- Global Ocean Forecasting
+- GLORYS12
+- FNO
+tasks: []
+datasets:
+- GLORYS12
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">GLONET</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+GLONET (Global Ocean Neural Network) is a global ocean neural-network forecasting system developed by Mercator Ocean International, a leading European ocean forecasting center.
+
+# Model Description
+
+GLONET forecasts global ocean states. It takes two consecutive daily states as input and outputs the 34-channel ocean state for the next day.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Global ocean forecast research | Train a dual-branch FNO/CNN ocean forecast model with GLORYS12-compatible data. |
+| Local quick validation | Use synthetic ocean fields to check data loading, pretraining, fine-tuning, inference, and visualization. |
+| ModelScope / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
+| Multi-GPU training | Run multi-GPU training 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/GLONET --local-dir ./GLONET
+cd GLONET
+```
+
+### 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 original work uses GLORYS12 reanalysis data. Real data must first be converted to the channel order and grid specified in `conf/config.yaml`; the raw GLORYS12 data is not included in this package. The default synthetic data is only for interface checks:
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Multi-GPU:
+
+```bash
+torchrun --nproc_per_node=8 scripts/train.py
+```
+
+Checkpoints are saved to `data/checkpoints/` by default.
+
+### Training Weights
+
+This repository provides weights trained on GLORYS12 data in the `weight/` folder. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+The prediction tensor is written to `result/glonet/data/prediction.pt` by default.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+The default output is `result/glonet/prediction.png`. Meaningful errors are computed only when a real reference field is provided.
+
+# 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 a reproduction of the original GLONET paper.

conf/config.yaml

@@ -0,0 +1,61 @@
+project:
+  name: glonet_reference
+  seed: 42
+  result_dir: ./result/glonet
+
+data:
+  data_dir: ./data/
+  channels:
+    - temperature_0.49m
+    - temperature_47m
+    - temperature_92m
+    - temperature_155m
+    - salinity_222m
+    - salinity_318m
+    - salinity_380m
+    - salinity_453m
+    - salinity_541m
+    - u_643m
+    - u_763m
+    - u_902m
+    - u_1245m
+    - u_1684m
+    - u_2225m
+    - u_3220m
+    - u_3597m
+    - u_3992m
+    - u_4405m
+    - u_4833m
+    - u_5274m
+    - v_643m
+    - v_763m
+    - v_902m
+    - v_1245m
+    - v_1684m
+    - v_2225m
+    - v_3220m
+    - v_3597m
+    - v_3992m
+    - v_4405m
+    - v_4833m
+    - v_5274m
+    - ssh_surface
+  grid: [16, 32]
+  input_steps: 2
+  output_steps: 4
+  synthetic_samples: 32
+  batch_size: 2
+
+model:
+  hidden_channels: 32
+  modes: [6, 8]
+  layers: 4
+
+training:
+  epochs: 10
+  learning_rate: 0.0001
+  device: auto
+  checkpoint_dir: ./data/checkpoints
+  checkpoint: ./data/checkpoints/model_glonet.pth
+  pretrain_rollout_steps: 1
+  finetune_rollout_steps: 4

config.json

@@ -0,0 +1,43 @@
+{
+  "model_name": "GLONET",
+  "model_type": "glonet",
+  "architectures": [
+    "GLONET"
+  ],
+  "framework": "PyTorch",
+  "domain": "ocean",
+  "task": "global-ocean-forecasting",
+  "implementation": {
+    "entry_point": "model/glonet.py",
+    "scope": "two-branch FNO/CNN reference implementation"
+  },
+  "architecture": {
+    "family": "Fourier neural operator with convolutional branch",
+    "grid_shape": [
+      16,
+      32
+    ],
+    "input_steps": 2,
+    "input_channels_per_step": 34,
+    "input_channels": 68,
+    "output_channels": 34,
+    "hidden_channels": 32,
+    "spectral_modes": [
+      6,
+      8
+    ],
+    "spectral_layers": 4,
+    "activation": "GELU"
+  },
+  "data": {
+    "dataset": "GLORYS12",
+    "time_step_days": 1,
+    "output_steps": 4,
+    "protocol": "synthetic_glorys12_shape"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/glonet.py",
+    "scripts/data_loader.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "ocean_forecasting",
+  "model": "GLONET",
+  "input_format": "BTCHW",
+  "protocol": "synthetic_glorys12_shape",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/glonet.py

@@ -0,0 +1,71 @@
+"""GLONET reference architecture based on the public paper description."""
+
+import torch
+from torch import nn
+
+
+class SpectralConv2d(nn.Module):
+    def __init__(self, channels, modes):
+        super().__init__()
+        self.modes_y, self.modes_x = modes
+        self.weight = nn.Parameter(torch.randn(channels, channels, self.modes_y, self.modes_x, 2) * 0.02)
+
+    def forward(self, x):
+        height, width = x.shape[-2:]
+        spectrum = torch.fft.rfft2(x, norm="ortho")
+        out = torch.zeros_like(spectrum)
+        modes_y = min(self.modes_y, height)
+        modes_x = min(self.modes_x, spectrum.shape[-1])
+        weight = torch.view_as_complex(self.weight[:, :, :modes_y, :modes_x].contiguous())
+        out[:, :, :modes_y, :modes_x] = torch.einsum(
+            "bixy,ioxy->boxy", spectrum[:, :, :modes_y, :modes_x], weight
+        )
+        return torch.fft.irfft2(out, s=(height, width), norm="ortho")
+
+
+class SpectralBlock(nn.Module):
+    def __init__(self, channels, modes):
+        super().__init__()
+        self.spectral = SpectralConv2d(channels, modes)
+        self.pointwise = nn.Conv2d(channels, channels, 1)
+        self.activation = nn.GELU()
+
+    def forward(self, x):
+        return self.activation(self.spectral(x) + self.pointwise(x))
+
+
+class CNNBranch(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(channels, channels, 3, padding=1), nn.GELU(),
+            nn.Conv2d(channels, channels, 3, padding=1), nn.GELU(),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class GLONET(nn.Module):
+    """Two-day to one-day global ocean forecast reference model.
+
+    The paper does not publish a complete layer configuration, so all sizing
+    choices remain explicit constructor parameters rather than hidden claims.
+    """
+
+    def __init__(self, in_channels, out_channels=None, hidden_channels=32, modes=(6, 8), layers=4):
+        super().__init__()
+        out_channels = out_channels or in_channels
+        self.input_projection = nn.Conv2d(in_channels, hidden_channels, 1)
+        self.fno = nn.Sequential(*[SpectralBlock(hidden_channels, modes) for _ in range(layers)])
+        self.cnn = CNNBranch(hidden_channels)
+        self.output_projection = nn.Sequential(
+            nn.Conv2d(hidden_channels * 2, hidden_channels, 1), nn.GELU(),
+            nn.Conv2d(hidden_channels, out_channels, 1),
+        )
+
+    def forward(self, x):
+        if x.ndim == 5:
+            x = x.flatten(1, 2)
+        features = self.input_projection(x)
+        return self.output_projection(torch.cat((self.fno(features), self.cnn(features)), dim=1))

scripts/data_loader.py

@@ -0,0 +1,31 @@
+"""Synthetic GLORYS12-shaped data for smoke tests and local development."""
+
+import glob
+import torch
+from torch.utils.data import Dataset
+
+try:
+    import h5py
+except ImportError:
+    h5py = None
+
+
+class SyntheticOceanDataset(Dataset):
+    def __init__(self, samples, channels, grid, input_steps=2, output_steps=1, seed=42, data_dir=None):
+        generator = torch.Generator().manual_seed(seed)
+        files = sorted(glob.glob(f"{data_dir}/data/*.h5")) if data_dir else []
+        if files and h5py is not None:
+            with h5py.File(files[0], "r") as handle:
+                fields = torch.from_numpy(handle["fields"][:]).float()
+            total = min(samples, fields.shape[0] - input_steps - output_steps + 1)
+            self.x = torch.stack([fields[i:i + input_steps] for i in range(total)])
+            self.y = torch.stack([fields[i + input_steps:i + input_steps + output_steps] for i in range(total)])
+        else:
+            self.x = torch.randn(samples, input_steps, channels, *grid, generator=generator)
+            self.y = torch.randn(samples, output_steps, channels, *grid, generator=generator)
+
+    def __len__(self):
+        return self.x.shape[0]
+
+    def __getitem__(self, index):
+        return self.x[index], self.y[index]

scripts/fake_data.py

@@ -0,0 +1,50 @@
+"""Generate a small synthetic dataset description for local smoke tests."""
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+try:
+    import h5py
+except ImportError as exc:
+    raise SystemExit("fake_data.py requires h5py; install it in the active OneScience environment") from exc
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
+    parser.add_argument("--output", default=None)
+    args = parser.parse_args()
+    with open(args.config, encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    metadata = {
+        "variables": config["data"]["channels"],
+        "grid": config["data"]["grid"],
+        "input_steps": config["data"]["input_steps"],
+        "output_steps": config["data"]["output_steps"],
+        "time_resolution": "1 day",
+        "source": "synthetic; not GLORYS12 values",
+    }
+    data_root = ROOT / config["data"]["data_dir"] / "data"
+    data_root.mkdir(parents=True, exist_ok=True)
+    fields = np.random.default_rng(config["project"]["seed"]).standard_normal(
+        (config["data"]["synthetic_samples"] + config["data"]["input_steps"],
+         len(config["data"]["channels"]), *config["data"]["grid"]), dtype=np.float32
+    )
+    with h5py.File(data_root / "2000.h5", "w") as handle:
+        dataset = handle.create_dataset("fields", data=fields)
+        dataset.attrs["variables"] = config["data"]["channels"]
+        dataset.attrs["time_step"] = 24
+    output = Path(args.output or ROOT / config["data"]["data_dir"] / "synthetic_metadata.json")
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
+    print(f"saved={output}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,43 @@
+import argparse
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+import torch
+import yaml
+
+from data_loader import SyntheticOceanDataset
+from model.glonet import GLONET
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
+    parser.add_argument("--checkpoint", default=None)
+    args = parser.parse_args()
+    with open(args.config, encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    channels = len(config["data"]["channels"])
+    model = GLONET(channels * config["data"]["input_steps"], out_channels=channels,
+                   hidden_channels=config["model"]["hidden_channels"], modes=config["model"]["modes"],
+                   layers=config["model"]["layers"])
+    checkpoint = Path(args.checkpoint or ROOT / config["training"]["checkpoint"])
+    state = torch.load(checkpoint, map_location="cpu", weights_only=False)
+    model.load_state_dict(state["model"])
+    model.eval()
+    sample, _ = SyntheticOceanDataset(1, channels, config["data"]["grid"],
+                                      input_steps=config["data"]["input_steps"],
+                                      output_steps=config["data"]["output_steps"],
+                                      data_dir=str(ROOT / config["data"]["data_dir"]))[0]
+    with torch.no_grad():
+        prediction = model(sample.unsqueeze(0))
+    output = ROOT / config["project"]["result_dir"] / "data" / "prediction.pt"
+    output.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(prediction, output)
+    print(f"prediction_shape={tuple(prediction.shape)}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,139 @@
+import argparse
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+import math
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import yaml
+
+
+def load_field(path, time_index=0):
+    """Return a prediction as [C, H, W] from common model output layouts."""
+    value = torch.load(path, map_location="cpu", weights_only=False)
+    if isinstance(value, dict):
+        for key in ("prediction", "predictions", "output", "outputs"):
+            if key in value:
+                value = value[key]
+                break
+    field = torch.as_tensor(value).detach().cpu().float().numpy()
+    if field.ndim == 5:  # [B, T, C, H, W]
+        field = field[0, time_index]
+    elif field.ndim == 4:  # [B, C, H, W] or [T, C, H, W]
+        field = field[0 if field.shape[0] == 1 else time_index]
+    elif field.ndim != 3:
+        raise ValueError(f"Expected [C,H,W], [B,C,H,W], or [B,T,C,H,W], got {field.shape}")
+    if field.ndim != 3:
+        raise ValueError(f"Selected output is not [C,H,W]: {field.shape}")
+    return field
+
+
+def load_channel_names(config_path, channel_count):
+    if config_path is None:
+        return [f"channel_{index}" for index in range(channel_count)]
+    with open(config_path, encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    names = config.get("data", {}).get("channels", [])
+    if len(names) != channel_count:
+        return [f"channel_{index}" for index in range(channel_count)]
+    return names
+
+
+def choose_channels(names, requested, max_panels):
+    if requested:
+        selected = []
+        for item in requested:
+            if item.isdigit():
+                index = int(item)
+                if not 0 <= index < len(names):
+                    raise ValueError(f"Channel index out of range: {index}")
+            else:
+                if item not in names:
+                    raise ValueError(f"Unknown channel: {item}")
+                index = names.index(item)
+            if index not in selected:
+                selected.append(index)
+        return selected
+    return list(range(min(max_panels, len(names))))
+
+
+def is_signed_channel(name):
+    return name.startswith(("u_", "v_")) or name.startswith("ssh_")
+
+
+def plot_fields(field, names, indices, output, title, reference=None):
+    columns = min(3, len(indices))
+    rows = math.ceil(len(indices) / columns)
+    has_reference = reference is not None
+    fig, axes = plt.subplots(rows, columns, figsize=(5.6 * columns, 4.4 * rows), squeeze=False)
+    axes = axes.ravel()
+    height, width = field.shape[-2:]
+    longitude = np.linspace(0, 360, width, endpoint=False)
+    latitude = np.linspace(90, -90, height)
+    extent = [longitude[0], longitude[-1], latitude[-1], latitude[0]]
+
+    for axis, index in zip(axes, indices):
+        data = field[index]
+        reference_data = reference[index] if has_reference else None
+        if reference_data is not None:
+            data_min = min(np.nanpercentile(data, 2), np.nanpercentile(reference_data, 2))
+            data_max = max(np.nanpercentile(data, 98), np.nanpercentile(reference_data, 98))
+        else:
+            data_min, data_max = np.nanpercentile(data, [2, 98])
+        if np.isclose(data_min, data_max):
+            data_min, data_max = float(np.nanmin(data)), float(np.nanmax(data) + 1e-6)
+        cmap = "RdBu_r" if is_signed_channel(names[index]) else "viridis"
+        image = axis.imshow(data, extent=extent, origin="upper", cmap=cmap,
+                            vmin=data_min, vmax=data_max, aspect="auto")
+        axis.set_title(names[index], fontsize=11, fontweight="bold")
+        axis.set_xlabel("Longitude (degrees)")
+        axis.set_ylabel("Latitude (degrees)")
+        axis.set_xticks([0, 90, 180, 270, 360])
+        axis.set_yticks([-90, -45, 0, 45, 90])
+        axis.grid(color="white", linewidth=0.35, alpha=0.35)
+        colorbar = fig.colorbar(image, ax=axis, fraction=0.046, pad=0.04)
+        colorbar.ax.tick_params(labelsize=8)
+        stats = f"min {np.nanmin(data):.3g} | max {np.nanmax(data):.3g} | mean {np.nanmean(data):.3g}"
+        if reference_data is not None:
+            rmse = np.sqrt(np.nanmean((data - reference_data) ** 2))
+            stats += f" | RMSE {rmse:.3g}"
+        axis.text(0.02, 0.02, stats, transform=axis.transAxes, fontsize=8,
+                  color="white", bbox={"facecolor": "black", "alpha": 0.55, "pad": 3})
+
+    for axis in axes[len(indices):]:
+        axis.remove()
+    fig.suptitle(title, fontsize=15, fontweight="bold")
+    fig.tight_layout()
+    fig.savefig(output, dpi=180, bbox_inches="tight")
+    plt.close(fig)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", default=str(ROOT / "result/glonet/data/prediction.pt"))
+    parser.add_argument("--output", default=str(ROOT / "result/glonet/prediction.png"))
+    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
+    parser.add_argument("--reference", default=None, help="Optional .pt truth field for RMSE comparison")
+    parser.add_argument("--channel", action="append", help="Channel name or zero-based index; repeatable")
+    parser.add_argument("--max-panels", type=int, default=6)
+    parser.add_argument("--time-index", type=int, default=0)
+    args = parser.parse_args()
+    prediction = load_field(args.input, args.time_index)
+    names = load_channel_names(args.config, prediction.shape[0])
+    indices = choose_channels(names, args.channel, args.max_panels)
+    reference = load_field(args.reference, args.time_index) if args.reference else None
+    if reference is not None and reference.shape != prediction.shape:
+        raise ValueError(f"Prediction/reference shape mismatch: {prediction.shape} vs {reference.shape}")
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+    plot_fields(prediction, names, indices, args.output,
+                f"GLONET ocean forecast | {len(indices)} channel(s)", reference)
+    print(f"saved={args.output}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,99 @@
+import argparse
+import os
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+import torch
+import torch.distributed as dist
+import yaml
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader, DistributedSampler
+
+from data_loader import SyntheticOceanDataset
+from model.glonet import GLONET
+
+
+def setup_distributed():
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    if world_size == 1:
+        return 0, 0, torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    local_rank = int(os.environ["LOCAL_RANK"])
+    force_cpu = os.environ.get("GLONET_FORCE_CPU", "0") == "1"
+    if torch.cuda.is_available() and not force_cpu:
+        device_count = torch.cuda.device_count()
+        if local_rank >= device_count:
+            raise RuntimeError(
+                f"LOCAL_RANK={local_rank} but only {device_count} accelerator(s) are visible; "
+                "reduce --nproc_per_node or fix CUDA_VISIBLE_DEVICES."
+            )
+        torch.cuda.set_device(local_rank)
+        device = torch.device("cuda", local_rank)
+        backend = "nccl"
+    else:
+        device = torch.device("cpu")
+        backend = "gloo"
+    dist.init_process_group(backend=backend, init_method="env://")
+    return dist.get_rank(), local_rank, device
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
+    parser.add_argument("--phase", choices=["pretrain", "finetune"], default="pretrain")
+    args = parser.parse_args()
+    with open(args.config, encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    rank, local_rank, device = setup_distributed()
+    torch.manual_seed(config["project"]["seed"] + rank)
+    channels = len(config["data"]["channels"])
+    rollout_steps = config["training"][f"{args.phase}_rollout_steps"]
+    dataset = SyntheticOceanDataset(config["data"]["synthetic_samples"], channels, config["data"]["grid"],
+                                    input_steps=config["data"]["input_steps"],
+                                    output_steps=config["data"]["output_steps"],
+                                    data_dir=str(ROOT / config["data"]["data_dir"]))
+    sampler = DistributedSampler(dataset, shuffle=True) if dist.is_initialized() else None
+    loader = DataLoader(dataset, batch_size=config["data"]["batch_size"], shuffle=sampler is None, sampler=sampler)
+    model = GLONET(channels * config["data"]["input_steps"], out_channels=channels,
+                   hidden_channels=config["model"]["hidden_channels"], modes=config["model"]["modes"],
+                   layers=config["model"]["layers"]).to(device)
+    checkpoint = ROOT / config["training"]["checkpoint"]
+    if args.phase == "finetune" and checkpoint.exists():
+        state = torch.load(checkpoint, map_location=device, weights_only=False)
+        model.load_state_dict(state["model"])
+    if dist.is_initialized():
+        model = DDP(model, device_ids=[local_rank] if device.type == "cuda" else None)
+    optimizer = torch.optim.Adam(model.parameters(), lr=config["training"]["learning_rate"])
+    for epoch in range(config["training"]["epochs"]):
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        model.train()
+        total = 0.0
+        for inputs, targets in loader:
+            inputs, targets = inputs.to(device), targets.to(device)
+            optimizer.zero_grad(set_to_none=True)
+            loss = 0.0
+            state = inputs
+            for step in range(rollout_steps):
+                prediction = model(state)
+                loss = loss + torch.nn.functional.mse_loss(prediction, targets[:, step])
+                state = torch.cat((state[:, 1:], prediction.unsqueeze(1)), dim=1)
+            loss = loss / rollout_steps
+            loss.backward()
+            optimizer.step()
+            total += loss.item()
+        if rank == 0:
+            print(f"epoch={epoch + 1} loss={total / len(loader):.6f}")
+    if rank == 0:
+        checkpoint.parent.mkdir(parents=True, exist_ok=True)
+        torch.save({"model": model.module.state_dict() if hasattr(model, "module") else model.state_dict(),
+                    "config": config}, checkpoint)
+        print(f"saved={checkpoint}")
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()