Upload folder using huggingface_hub

README.md

@@ -0,0 +1,140 @@
+---
+license: apache-2.0
+language:
+- en
+- zh
+tags:
+- OneScience
+- Earth Science
+- Weather Forecast
+- Short-to-Medium-Range Weather Forecast
+- ERA5
+frameworks: PyTorch
+datasets:
+  - OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">FengWu</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+FengWu is a global medium-range weather forecast foundation model jointly released by the Shanghai Artificial Intelligence Laboratory and multiple universities. It has been adopted by organizations such as the Hong Kong Observatory for operational weather forecasting.
+
+Paper: FengWu: Pushing the Skillful Global Medium-range Weather Forecast beyond 10 Days Lead
+
+https://arxiv.org/abs/2304.02948
+
+# Model Description
+
+The FengWu model is built on a multi-modal and multi-task deep learning approach, without relying on traditional physical equations. It is trained entirely on ERA5 reanalysis data.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Weather Forecast Training | Train FengWu using ERA5 HDF5 data |
+| Local Quick Validation | Use synthetic data to verify data loading, model training & inference, and inference result visualization. |
+| ModelScope / OneCode Execution | Download as a standalone model package, install dependencies, and run scripts directly. |
+| Multi-GPU Training | Launch multi-process training via `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
+modelscope download --model OneScience/FengWu --local_dir ./FengWu
+cd FengWu
+```
+
+### 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 ERA5 data for training (due to file size limits, the current repository contains a slice of the full dataset). Users can download it with the command below and confirm that the data path in `conf/config.yaml` is set correctly:
+
+```bash
+modelscope download --dataset OneScience/ERA5 --local_dir ./data
+```
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Multi-GPU:
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+Training will save `model_bak.pth` under `data/checkpoints/`.
+
+### Training Weights
+
+This repository provides weights trained on 39 years of ERA5 reanalysis data in the `weights/` folder. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference results will be saved to `result/output/`.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+You can specify a date and variable at the end of `result.py` for visualization.
+
+# OneScience Official Information
+
+| 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 & License
+
+- This repository is a reproduction of the original FengWu paper.

conf/config.yaml

@@ -0,0 +1,145 @@
+model:
+
+    start_epoch: 0
+    max_epoch: 1
+    lr: 1E-3
+    num_blocks: 8
+    patch_size: [4, 4]
+    embed_dim: 192
+    num_heads: [6, 12, 12, 6]
+    window_size: [2, 6, 12]
+    pressure_level: 37
+    checkpoint_dir: "./data/checkpoints"
+    patience: 50
+
+ # 整个数据读取流程
+datapipe:
+  name: "ERA5"
+  task: "weather_forecasting"
+
+  # dataset设定
+  dataset:
+    type: "hdf5"
+    data_dir: './data/' # "$ONESCIENCE_DATASETS_DIR/ERA5/newh5/" 
+    
+    train_time: [2000, 2001]
+    val_time: [2002]
+    test_time: [2003]
+    img_size: [721, 1440]
+    verbose: true
+    cache: false
+    # 气象变量
+
+    channels:  ['10m_u_component_of_wind', '10m_v_component_of_wind', '2m_temperature', 'mean_sea_level_pressure', 
+    
+                'geopotential_1', 'geopotential_2', 'geopotential_3', 'geopotential_5', 'geopotential_7', 
+                'geopotential_10', 'geopotential_20', 'geopotential_30', 'geopotential_50', 'geopotential_70', 
+                'geopotential_100', 'geopotential_125', 'geopotential_150', 'geopotential_175', 'geopotential_200', 
+                'geopotential_225', 'geopotential_250', 'geopotential_300', 'geopotential_350', 'geopotential_400', 
+                'geopotential_450', 'geopotential_500', 'geopotential_550', 'geopotential_600', 'geopotential_650', 
+                'geopotential_700', 'geopotential_750', 'geopotential_775', 'geopotential_800', 'geopotential_825', 
+                'geopotential_850', 'geopotential_875', 'geopotential_900', 'geopotential_925', 'geopotential_950', 
+                'geopotential_975', 'geopotential_1000',
+
+                'relative_humidity_1', 'relative_humidity_2', 'relative_humidity_3', 'relative_humidity_5', 'relative_humidity_7', 
+                'relative_humidity_10', 'relative_humidity_20', 'relative_humidity_30', 'relative_humidity_50', 'relative_humidity_70', 
+                'relative_humidity_100', 'relative_humidity_125', 'relative_humidity_150', 'relative_humidity_175', 'relative_humidity_200', 
+                'relative_humidity_225', 'relative_humidity_250', 'relative_humidity_300', 'relative_humidity_350', 'relative_humidity_400', 
+                'relative_humidity_450', 'relative_humidity_500', 'relative_humidity_550', 'relative_humidity_600', 'relative_humidity_650', 
+                'relative_humidity_700', 'relative_humidity_750', 'relative_humidity_775', 'relative_humidity_800', 'relative_humidity_825', 
+                'relative_humidity_850', 'relative_humidity_875', 'relative_humidity_900', 'relative_humidity_925', 'relative_humidity_950', 
+                'relative_humidity_975', 'relative_humidity_1000',
+
+                'u_component_of_wind_1', 'u_component_of_wind_2', 'u_component_of_wind_3', 'u_component_of_wind_5', 'u_component_of_wind_7', 
+                'u_component_of_wind_10', 'u_component_of_wind_20', 'u_component_of_wind_30', 'u_component_of_wind_50', 'u_component_of_wind_70', 
+                'u_component_of_wind_100', 'u_component_of_wind_125', 'u_component_of_wind_150', 'u_component_of_wind_175', 'u_component_of_wind_200', 
+                'u_component_of_wind_225', 'u_component_of_wind_250', 'u_component_of_wind_300', 'u_component_of_wind_350', 'u_component_of_wind_400', 
+                'u_component_of_wind_450', 'u_component_of_wind_500', 'u_component_of_wind_550', 'u_component_of_wind_600', 'u_component_of_wind_650', 
+                'u_component_of_wind_700', 'u_component_of_wind_750', 'u_component_of_wind_775', 'u_component_of_wind_800', 'u_component_of_wind_825', 
+                'u_component_of_wind_850', 'u_component_of_wind_875', 'u_component_of_wind_900', 'u_component_of_wind_925', 'u_component_of_wind_950', 
+                'u_component_of_wind_975', 'u_component_of_wind_1000',
+
+                'v_component_of_wind_1', 'v_component_of_wind_2', 'v_component_of_wind_3', 'v_component_of_wind_5', 'v_component_of_wind_7', 
+                'v_component_of_wind_10', 'v_component_of_wind_20', 'v_component_of_wind_30', 'v_component_of_wind_50', 'v_component_of_wind_70', 
+                'v_component_of_wind_100', 'v_component_of_wind_125', 'v_component_of_wind_150', 'v_component_of_wind_175', 'v_component_of_wind_200', 
+                'v_component_of_wind_225', 'v_component_of_wind_250', 'v_component_of_wind_300', 'v_component_of_wind_350', 'v_component_of_wind_400', 
+                'v_component_of_wind_450', 'v_component_of_wind_500', 'v_component_of_wind_550', 'v_component_of_wind_600', 'v_component_of_wind_650', 
+                'v_component_of_wind_700', 'v_component_of_wind_750', 'v_component_of_wind_775', 'v_component_of_wind_800', 'v_component_of_wind_825', 
+                'v_component_of_wind_850', 'v_component_of_wind_875', 'v_component_of_wind_900', 'v_component_of_wind_925', 'v_component_of_wind_950', 
+                'v_component_of_wind_975', 'v_component_of_wind_1000',
+
+                'temperature_1', 'temperature_2', 'temperature_3', 'temperature_5', 'temperature_7', 
+                'temperature_10', 'temperature_20', 'temperature_30', 'temperature_50', 'temperature_70', 
+                'temperature_100', 'temperature_125', 'temperature_150', 'temperature_175', 'temperature_200', 
+                'temperature_225', 'temperature_250', 'temperature_300', 'temperature_350', 'temperature_400', 
+                'temperature_450', 'temperature_500', 'temperature_550', 'temperature_600', 'temperature_650', 
+                'temperature_700', 'temperature_750', 'temperature_775', 'temperature_800', 'temperature_825', 
+                'temperature_850', 'temperature_875', 'temperature_900', 'temperature_925', 'temperature_950', 
+                'temperature_975', 'temperature_1000',
+                ]
+
+    variables:
+      - "u10"        # 10m U wind component
+      - "v10"        # 10m V wind component
+      - "t2m"        # 2m temperature
+      - "msl"        # Mean sea level pressure
+      - "z500"       # Geopotential at 500 hPa
+      - "t850"       # Temperature at 850 hPa
+    
+    # 时间配置
+    time_range: ["2000-01-01", "2020-12-31"]
+    time_steps: 1
+    time_res: 6
+    
+    # 空间配置
+    spatial_resolution: [0.25, 0.25]
+    
+    # 采样配置
+    num_samples: -1  # -1 表示使用全部数据
+    shuffle: true
+    random_seed: 42
+    
+    # 领域特定配置
+    extra:
+      levels: [500, 850, 1000]
+      lat_range: [-90, 90]
+      lon_range: [0, 360]
+  
+  # 数据转换配置
+  transforms:
+    - type: "Normalize"
+      params:
+        mean: [0.0, 0.0, 288.0, 101325.0, 50000.0, 270.0]
+        std: [5.0, 5.0, 15.0, 1000.0, 5000.0, 10.0]
+        keys: ["input", "target"]
+    
+    - type: "ToTensor"
+      params:
+        keys: null  # null表示转换所有numpy数组
+  
+  # 其他配置
+
+
+  # DataLoader配置
+  dataloader:
+    mask_dtype: "float32"
+    batch_size: 1
+    num_workers: 1
+    pin_memory: true
+    drop_last: true
+    shuffle: false  # 使用sampler时设为false
+    prefetch_factor: 2
+    persistent_workers: true
+  
+  # 分布式配置
+  distributed:
+    enabled: true
+    sampler: "DistributedSampler"
+    rank: 0
+    world_size: 4
+    shuffle: true
+    seed: 42
+    drop_last: true
+
+
+

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"other"}

model/fengwu.py

@@ -0,0 +1,278 @@
+import math
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+from torch import nn
+from onescience.modules.encoder.fengwuencoder import FengWuEncoder
+from onescience.modules.decoder.fengwudecoder import FengWuDecoder
+from onescience.modules.fuser.fengwufuser import FengWuFuser
+
+from onescience.models.meta import ModelMetaData
+
+
+@dataclass
+class MetaData(ModelMetaData):
+    name: str = "Fengwu"
+    # Optimization
+    jit: bool = False  # ONNX Ops Conflict
+    cuda_graphs: bool = True
+    amp: bool = True
+    # Inference
+    onnx_cpu: bool = False  # No FFT op on CPU
+    onnx_gpu: bool = True
+    onnx_runtime: bool = True
+    # Physics informed
+    var_dim: int = 1
+    func_torch: bool = False
+    auto_grad: bool = False
+
+
+class Fengwu(nn.Module):
+    """
+    FengWu 的主模型实现。
+
+    该模型由多个变量分支编码器、一个中分辨率三维 fuser，以及多个变量分支解码器组成。
+
+    结构顺序为：
+
+    - 多个 `FengWuEncoder`
+      - 分别编码 surface、Z、R、U、V、T 六个变量分支
+    - `FengWuFuser`
+      - 在统一三维网格 `(Variables, Height, Width)` 上融合中分辨率特征
+    - 多个 `FengWuDecoder`
+      - 分别恢复各变量分支输出
+
+    与 Pangu 不同，FengWu 不把所有变量直接拼成单一路径输入，而是先按变量族分支编码，
+    再在中分辨率层面做跨变量三维融合。
+
+    Reference:
+    - `FengWu: Pushing the Skillful Global Medium-range Weather Forecast beyond 10 Days Lead`
+    - https://arxiv.org/pdf/2304.02948.pdf
+
+    Args:
+        img_size (tuple[int, int]):
+            输入场空间尺寸 `(Height, Width)`。
+        pressure_level (int):
+            高空变量的层数。
+        embed_dim (int):
+            编码器高分辨率特征维度。
+        patch_size (tuple[int, int]):
+            二维 patch 切分尺寸 `(PatchHeight, PatchWidth)`。
+        num_heads (tuple[int, int, int, int]):
+            注意力头数配置，约定顺序为：
+            `(EncoderHighHeads, EncoderMiddleHeads, DecoderMiddleHeads, DecoderHighHeads)`。
+        window_size (tuple[int, int, int]):
+            FengWuFuser 的三维窗口大小。
+    """
+
+    def __init__(
+        self,
+        img_size=(721, 1440),
+        pressure_level=37,
+        embed_dim=192,
+        patch_size=(4, 4),
+        num_heads=(6, 12, 12, 6),
+        window_size=(2, 6, 12),
+    ):
+        super().__init__()
+        input_resolution = (
+            math.ceil(img_size[0] / patch_size[0]),
+            math.ceil(img_size[1] / patch_size[1]),
+        )
+        middle_resolution = (
+            math.ceil(input_resolution[0] / 2),
+            math.ceil(input_resolution[1] / 2),
+        )
+        encoder_num_heads = (num_heads[0], num_heads[1])
+        decoder_num_heads = (num_heads[3], num_heads[2])
+        drop_path = np.linspace(0, 0.2, 8).tolist()
+
+        self.encoder_surface = FengWuEncoder(
+            input_resolution=input_resolution,
+            middle_resolution=middle_resolution,
+            in_chans=4,
+            img_size=img_size,
+            patch_size=patch_size,
+            dim=embed_dim,
+            num_heads=encoder_num_heads,
+            window_size=window_size[1:],
+            drop_path=drop_path,
+        )
+        self.encoder_z = FengWuEncoder(
+            input_resolution=input_resolution,
+            middle_resolution=middle_resolution,
+            in_chans=pressure_level,
+            img_size=img_size,
+            patch_size=patch_size,
+            dim=embed_dim,
+            num_heads=encoder_num_heads,
+            window_size=window_size[1:],
+            drop_path=drop_path,
+        )
+        self.encoder_r = FengWuEncoder(
+            input_resolution=input_resolution,
+            middle_resolution=middle_resolution,
+            in_chans=pressure_level,
+            img_size=img_size,
+            patch_size=patch_size,
+            dim=embed_dim,
+            num_heads=encoder_num_heads,
+            window_size=window_size[1:],
+            drop_path=drop_path,
+        )
+        self.encoder_u = FengWuEncoder(
+            input_resolution=input_resolution,
+            middle_resolution=middle_resolution,
+            in_chans=pressure_level,
+            img_size=img_size,
+            patch_size=patch_size,
+            dim=embed_dim,
+            num_heads=encoder_num_heads,
+            window_size=window_size[1:],
+            drop_path=drop_path,
+        )
+        self.encoder_v = FengWuEncoder(
+            input_resolution=input_resolution,
+            middle_resolution=middle_resolution,
+            in_chans=pressure_level,
+            img_size=img_size,
+            patch_size=patch_size,
+            dim=embed_dim,
+            num_heads=encoder_num_heads,
+            window_size=window_size[1:],
+            drop_path=drop_path,
+        )
+        self.encoder_t = FengWuEncoder(
+            input_resolution=input_resolution,
+            middle_resolution=middle_resolution,
+            in_chans=pressure_level,
+            img_size=img_size,
+            patch_size=patch_size,
+            dim=embed_dim,
+            num_heads=encoder_num_heads,
+            window_size=window_size[1:],
+            drop_path=drop_path,
+        )
+
+        self.fuser = FengWuFuser(
+            input_resolution=(6, *middle_resolution),
+            dim=embed_dim * 2,
+            num_heads=num_heads[2],
+            window_size=window_size,
+            drop_path=drop_path[2:],
+        )
+
+        self.decoder_surface = FengWuDecoder(
+            output_resolution=input_resolution,
+            middle_resolution=middle_resolution,
+            out_chans=4,
+            img_size=img_size,
+            patch_size=patch_size,
+            dim=embed_dim,
+            num_heads=decoder_num_heads,
+            window_size=window_size[1:],
+            drop_path=drop_path,
+        )
+        self.decoder_z = FengWuDecoder(
+            output_resolution=input_resolution,
+            middle_resolution=middle_resolution,
+            out_chans=pressure_level,
+            img_size=img_size,
+            patch_size=patch_size,
+            dim=embed_dim,
+            num_heads=decoder_num_heads,
+            window_size=window_size[1:],
+            drop_path=drop_path,
+        )
+        self.decoder_r = FengWuDecoder(
+            output_resolution=input_resolution,
+            middle_resolution=middle_resolution,
+            out_chans=pressure_level,
+            img_size=img_size,
+            patch_size=patch_size,
+            dim=embed_dim,
+            num_heads=decoder_num_heads,
+            window_size=window_size[1:],
+            drop_path=drop_path,
+        )
+        self.decoder_u = FengWuDecoder(
+            output_resolution=input_resolution,
+            middle_resolution=middle_resolution,
+            out_chans=pressure_level,
+            img_size=img_size,
+            patch_size=patch_size,
+            dim=embed_dim,
+            num_heads=decoder_num_heads,
+            window_size=window_size[1:],
+            drop_path=drop_path,
+        )
+        self.decoder_v = FengWuDecoder(
+            output_resolution=input_resolution,
+            middle_resolution=middle_resolution,
+            out_chans=pressure_level,
+            img_size=img_size,
+            patch_size=patch_size,
+            dim=embed_dim,
+            num_heads=decoder_num_heads,
+            window_size=window_size[1:],
+            drop_path=drop_path,
+        )
+        self.decoder_t = FengWuDecoder(
+            output_resolution=input_resolution,
+            middle_resolution=middle_resolution,
+            out_chans=pressure_level,
+            img_size=img_size,
+            patch_size=patch_size,
+            dim=embed_dim,
+            num_heads=decoder_num_heads,
+            window_size=window_size[1:],
+            drop_path=drop_path,
+        )
+
+        self.img_size = img_size
+        self.pressure_level = pressure_level
+        self.patch_size = patch_size
+        self.input_resolution = input_resolution
+        self.middle_resolution = middle_resolution
+
+    def forward(self, surface, z, r, u, v, t):
+        surface, skip_surface = self.encoder_surface(surface)
+        z, skip_z = self.encoder_z(z)
+        r, skip_r = self.encoder_r(r)
+        u, skip_u = self.encoder_u(u)
+        v, skip_v = self.encoder_v(v)
+        t, skip_t = self.encoder_t(t)
+
+        x = torch.concat(
+            [
+                surface.unsqueeze(1),
+                z.unsqueeze(1),
+                r.unsqueeze(1),
+                u.unsqueeze(1),
+                v.unsqueeze(1),
+                t.unsqueeze(1),
+            ],
+            dim=1,
+        )
+        Batch, Variables, NumTokensPerVariable, Channels = x.shape
+        x = x.reshape(Batch, -1, Channels)
+        x = self.fuser(x)
+
+        x = x.reshape(Batch, Variables, NumTokensPerVariable, Channels)
+        surface, z, r, u, v, t = (
+            x[:, 0, :, :],
+            x[:, 1, :, :],
+            x[:, 2, :, :],
+            x[:, 3, :, :],
+            x[:, 4, :, :],
+            x[:, 5, :, :],
+        )
+
+        surface = self.decoder_surface([surface, skip_surface])
+        z = self.decoder_z([z, skip_z])
+        r = self.decoder_r([r, skip_r])
+        u = self.decoder_u([u, skip_u])
+        v = self.decoder_v([v, skip_v])
+        t = self.decoder_t([t, skip_t])
+        return surface, z, r, u, v, t

scripts/fake_data.py

@@ -0,0 +1,99 @@
+import os
+import sys
+from pathlib import Path
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+import h5py
+import numpy as np
+import xarray as xr
+from onescience.utils.YParams import YParams
+
+
+# 各数据集固定的空间和时间维度
+DATASET_DIMS = {"T": 10, "H": 721, "W": 1440, "time_step": 6}
+
+
+def generate_fake_h5(data_dir, var_names, years, dims):
+    """
+    为每个年份生成一个空 h5 文件。
+    利用 HDF5 chunked 数据集未写入 chunk 即返回 fill_value=0 的特性，
+    文件实际只含元数据，极小，但 shape 与真实数据完全一致。
+    均值/标准差也作为数据集内嵌进每年的 h5，与 era5.py 新版读取方式对应。
+    """
+    os.makedirs(os.path.join(data_dir, "data"), exist_ok=True)
+    T, C = dims["T"], len(var_names)
+    H, W = dims["H"], dims["W"]
+
+    means = np.zeros((1, C, 1, 1), dtype=np.float32)
+    stds  = np.ones((1, C, 1, 1), dtype=np.float32)
+
+    for year in years:
+        path = os.path.join(data_dir, "data", f"{year}.h5")
+        with h5py.File(path, "w") as f:
+            ds = f.create_dataset(
+                "fields",
+                shape=(T, C, H, W),
+                dtype="float32",
+                chunks=(1, C, H, W),
+                fillvalue=0.0,
+            )
+            ds.attrs["variables"] = var_names
+            ds.attrs["time_step"] = dims["time_step"]
+            f.create_dataset("global_means", data=means)
+            f.create_dataset("global_stds",  data=stds)
+
+        size_kb = os.path.getsize(path) / 1024
+        print(f"  {year}.h5  shape=({T},{C},{H},{W})  "
+              f"logical={T*C*H*W*4/1024**3:.1f}GB  actual={size_kb:.1f}KB")
+
+
+def get_static(data_dir, var, name):
+    os.makedirs(data_dir, exist_ok=True)
+    ds = xr.Dataset(
+        data_vars={
+            f"{var}": (("valid_time", "latitude", "longitude"),
+                np.random.rand(1, 721, 1440).astype(np.float32))
+        },
+        coords={
+            "valid_time": ["2015-12-31"],
+            "latitude": np.linspace(90, -90, 721, dtype=np.float64),
+            "longitude": np.linspace(0, 359.75, 1440, dtype=np.float64),
+            "number": 0,
+            "expver": "",
+        },
+        attrs={
+            "GRIB_centre": "ecmf",
+            "GRIB_centreDescription": "European Centre for Medium-Range Weather Forecasts",
+            "GRIB_subCentre": "0",
+            "Conventions": "CF-1.7",
+            "institution": "European Centre for Medium-Range Weather Forecasts",
+            "history": "Generated manually",
+        }
+    )
+
+    ds.to_netcdf(f"{data_dir}/{name}.nc")
+    arr = np.random.randn(721, 1440).astype(np.float32)
+    np.save(f'{data_dir}/land_mask.npy', arr)
+    np.save(f'{data_dir}/soil_type.npy', arr)
+    np.save(f'{data_dir}/topography.npy', arr)
+    print(f"✅ Static data: {arr.shape}, dtype: {arr.dtype}, save to {data_dir}")
+
+
+if __name__ == "__main__":
+    cfg_datapipe = YParams("conf/config.yaml", "datapipe")
+
+    if cfg_datapipe.dataset.data_dir.startswith("/public/") or cfg_datapipe.dataset.data_dir.startswith("/work2/"):
+        print("请检查 config，确保各 *_dir 指向本地测试路径而非生产路径。")
+        exit()
+
+    years    = cfg_datapipe.dataset.train_time + cfg_datapipe.dataset.val_time + cfg_datapipe.dataset.test_time
+    atm_vars = cfg_datapipe.dataset.channels
+
+    generate_fake_h5(cfg_datapipe.dataset.data_dir, atm_vars, years, DATASET_DIMS)
+
+    static_dir = os.path.join(cfg_datapipe.dataset.data_dir, "static")
+    get_static(static_dir, 'z', 'geopotential')
+    get_static(static_dir, 'lsm', 'land_sea_mask')
+
+
+    print("\n✅ Fake datasets generated.")

scripts/inference.py

@@ -0,0 +1,81 @@
+import torch
+import os
+import sys
+from pathlib import Path
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+import glob
+import numpy as np
+import h5py
+from tqdm import tqdm
+from model.fengwu import Fengwu
+from onescience.utils.YParams import YParams
+from onescience.datapipes.climate import ERA5Datapipe
+
+
+def get_stats(data_dir, channels):
+    """从新版 h5 中读取变量列表与归一化参数（均值/标准差）"""
+    h5_files = sorted(glob.glob(os.path.join(data_dir, "data", "*.h5")))
+    with h5py.File(h5_files[0], "r") as f:
+        ds = f["fields"]
+        all_variables = [v.decode() if isinstance(v, bytes) else v for v in ds.attrs["variables"]]
+        mu = f["global_means"][:]   # [1, C, 1, 1]
+        std = f["global_stds"][:]
+
+    channel_indices = [all_variables.index(v) for v in channels]
+    means = mu[:, channel_indices, :, :]
+    stds = std[:, channel_indices, :, :]
+    return means, stds
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+
+    ## Model config init
+    config_file_path = os.path.join(current_path, "conf/config.yaml")
+    cfg = YParams(config_file_path, "model")
+
+    ## DataLoader init
+    cfg_data = YParams(config_file_path, "datapipe")
+    means, stds = get_stats(cfg_data.dataset.data_dir, cfg_data.dataset.channels)
+
+    datapipe = ERA5Datapipe(
+        dataset_dir=cfg_data.dataset.data_dir,
+        used_variables=cfg_data.dataset.channels,
+        used_years=cfg_data.dataset.test_time,
+        distributed=False,
+        batch_size=1,
+        num_workers=4,
+    )
+    test_dataloader, _ = datapipe.get_dataloader("test")
+
+    ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location="cuda:0")
+    model = Fengwu(img_size=cfg_data.dataset.img_size,
+                   pressure_level=cfg.pressure_level,
+                   embed_dim=cfg.embed_dim,
+                   patch_size=cfg.patch_size,
+                   num_heads=cfg.num_heads,
+                   window_size=cfg.window_size,
+                   ).to('cuda:0')
+    model.load_state_dict(ckpt["model_state_dict"])
+
+    model.eval()
+    os.makedirs('result/output/', exist_ok=True)
+    print(f"📂 samples will be generated to './result/output/'")
+    with torch.no_grad():
+        for data in tqdm(test_dataloader, desc="Inferring testset", unit="batch"):
+            invar = data[0].to("cuda:0", dtype=torch.float32)
+            outvar = data[1].to("cuda:0", dtype=torch.float32)
+            filename = data[4][-1][0]
+            surface = invar[:, :4, :, :]
+            z = invar[:, 4:41, :, :]
+            r = invar[:, 41:78, :, :]
+            u = invar[:, 78:115, :, :]
+            v = invar[:, 115:152, :, :]
+            t = invar[:, 152:189, :, :]
+
+            surface_p, z_p, r_p, u_p, v_p, t_p = model(surface, z, r, u, v, t)
+            pred_var = torch.concat([surface_p, z_p, r_p, u_p, v_p, t_p], dim=1).cpu().numpy()
+            pred_var = pred_var * stds + means
+            np.save(f"result/output/{filename}.npy", pred_var)

scripts/result.py

@@ -0,0 +1,217 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import os
+import sys
+from pathlib import Path
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+import glob
+import h5py
+from datetime import datetime
+from tqdm import tqdm
+from onescience.utils.fcn.YParams import YParams
+from matplotlib import rcParams
+
+# rcParams['font.family'] = 'serif'
+# rcParams['font.serif'] = ['DejaVu Serif']
+rcParams['mathtext.fontset'] = 'stix'
+rcParams['axes.linewidth'] = 0.9
+rcParams['xtick.major.width'] = 0.9
+rcParams['ytick.major.width'] = 0.9
+
+
+def get_metadata(data_dir, channels):
+    """从新版 h5 attrs 中读取变量列表和 time_step"""
+    h5_files = sorted(glob.glob(os.path.join(data_dir, "data", "*.h5")))
+    with h5py.File(h5_files[0], "r") as f:
+        ds = f["fields"]
+        all_variables = [v.decode() if isinstance(v, bytes) else v for v in ds.attrs["variables"]]
+        time_step = int(ds.attrs["time_step"])
+
+    channel_indices = [all_variables.index(v) for v in channels]
+
+    total_files = [f for f in os.listdir('./result/output/') if f.endswith('.npy')]
+    total_files.sort()
+    return total_files, channel_indices, time_step
+
+
+def filename_to_index(filename, time_step):
+    """将 YYYYMMDDHH 格式的文件名转换为年度 h5 文件中的时间步索引"""
+    dt = datetime.strptime(filename, "%Y%m%d%H")
+    year_start = datetime(dt.year, 1, 1)
+    hours = (dt - year_start).total_seconds() / 3600
+    return int(hours / time_step)
+
+
+def get_result(total_files, channel_indices, time_step, data_dir, clim_mean):
+    channel_rmse = np.zeros(len(channel_indices))
+    channel_acc = np.zeros(len(channel_indices))
+    clim_mean = clim_mean[0, :, :, :]
+    if not os.path.exists('./result/rmse.npy') or not os.path.exists('result/acc.npy'):
+        numerator = np.zeros(len(channel_indices))
+        pred_sq_sum = np.zeros(len(channel_indices))
+        label_sq_sum = np.zeros(len(channel_indices))
+        for file in tqdm(total_files, unit="files"):
+            fname = file[:-4]  # 去掉 .npy
+            year = fname[:4]
+            t_idx = filename_to_index(fname, time_step)
+            with h5py.File(os.path.join(data_dir, 'data', f'{year}.h5'), "r") as f:
+                label = f["fields"][t_idx]  # [C, H, W]
+                label = label[channel_indices]
+            pred = np.load(f'result/output/{file}').squeeze()
+
+            label_anom = label - clim_mean
+            pred_anom = pred - clim_mean
+            # 累加
+            numerator += np.sum(pred_anom * label_anom, axis=(1, 2))
+            pred_sq_sum += np.sum(pred_anom ** 2, axis=(1, 2))
+            label_sq_sum += np.sum(label_anom ** 2, axis=(1, 2))
+
+            channel_rmse += np.sqrt(np.mean((label - pred) ** 2, axis=(1, 2)))
+        channel_rmse /= len(total_files)
+        channel_acc = numerator / (np.sqrt(pred_sq_sum * label_sq_sum) + 1e-8)
+        np.save('./result/acc.npy', channel_acc)
+        np.save('./result/rmse.npy', channel_rmse)
+
+
+def show_result():
+    channel_rmse = np.load('./result/rmse.npy')
+    channel_acc = np.load('./result/acc.npy')
+
+    channels = [cfg_data.dataset.channels[i] for i in range(len(channel_indices))]
+    w = 24  # 最长 channel 名宽度
+
+    # 表头
+    print(f"┌{'─' * (w + 2)}┬{'─' * 14}┬{'─' * 14}┐")
+    print(f"│ {'Channel':<{w}} │ {'RMSE':>12} │ {'ACC':>12} │")
+    print(f"├{'─' * (w + 2)}┼{'─' * 14}┼{'─' * 14}┤")
+    # 数据行
+    for i, ch in enumerate(channels):
+        print(f"│ {ch:<{w}} │ {channel_rmse[i]:>12.4f} | {channel_acc[i]:>12.4f} |")
+    print(f"├{'─' * (w + 2)}┼{'─' * 14}┼{'─' * 14}┤")
+    print(f"│ {'Average':<{w}} │ {np.mean(channel_rmse):>12.4f} │ {np.mean(channel_acc):>12.4f} │")
+    print(f"└{'─' * (w + 2)}┴{'─' * 14}┴{'─' * 14}┘")
+
+
+def plot(label, pred, var, filename):
+    # 基础设置
+    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
+
+    # 坐标轴标签
+    xtick_labels = ['180°W', '90°W', '0°', '90°E', '180°E']
+    ytick_labels = ['90°S', '45°S', '0°', '45°N', '90°N']
+    xticks = np.linspace(0, label.shape[-1] - 1, 5)
+    yticks = np.linspace(0, label.shape[-2] - 1, 5)
+
+    # 计算统一色条范围
+    vmin = min(label.min(), pred.min())
+    vmax = max(label.max(), pred.max())
+
+    # 计算差异和 RMSE
+    diff = label - pred
+    rmse = np.sqrt(np.mean(diff ** 2))
+    diff_abs_max = np.abs(diff).max()
+
+    # 绘图配置
+    plot_configs = [
+        {'data': label, 'title': 'Truth', 'cmap': 'viridis', 'vmin': vmin, 'vmax': vmax},
+        {'data': pred,  'title': 'Prediction', 'cmap': 'viridis', 'vmin': vmin, 'vmax': vmax},
+        {'data': diff,  'title': f'Difference (RMSE={rmse:.2f})', 'cmap': 'RdBu_r', 'vmin': -diff_abs_max, 'vmax': diff_abs_max},
+    ]
+
+    # 统一绘制
+    for ax, cfg in zip(axes, plot_configs):
+        im = ax.imshow(cfg['data'], cmap=cfg['cmap'], vmin=cfg['vmin'], vmax=cfg['vmax'])
+        ax.set_title(cfg['title'], fontsize=12, pad=4)
+        ax.set_xlabel('Longitude')
+        ax.set_ylabel('Latitude')
+        ax.set_xticks(xticks)
+        ax.set_xticklabels(xtick_labels)
+        ax.set_yticks(yticks)
+        ax.set_yticklabels(ytick_labels)
+        plt.colorbar(im, ax=ax, orientation='horizontal')
+
+    # 总标题
+    fig.suptitle(var, fontsize=14, fontweight='bold', y=0.98)
+
+    plt.savefig(filename, dpi=300, bbox_inches='tight')
+    plt.close()
+
+
+def plot_loss(train_loss, valid_loss):
+
+    mask = ~(np.isnan(train_loss) | np.isnan(valid_loss))
+    train_loss = train_loss[mask]
+    valid_loss = valid_loss[mask]
+
+    fig, ax = plt.subplots(figsize=(5, 3.5))
+    # 配置
+    colors = {'train': '#2563EB', 'valid': '#EA580C'}
+    epochs = np.arange(1, len(train_loss) + 1)
+
+    # 绑定曲线
+    ax.plot(epochs, train_loss, color=colors['train'], linewidth=1.5, label='Train')
+    ax.plot(epochs, valid_loss, color=colors['valid'], linewidth=1.5, label='Valid', linestyle='--')
+    # 标注最小值
+    min_idx = np.argmin(valid_loss)
+    ax.scatter(epochs[min_idx], valid_loss[min_idx],
+               color=colors['valid'], s=40, zorder=5, edgecolors='white')
+    ax.annotate(f'Best: {valid_loss[min_idx]:.3f}',
+                xy=(epochs[min_idx], valid_loss[min_idx]),
+                xytext=(10, 10), textcoords='offset points', fontsize=8, color=colors['valid'],
+                arrowprops=dict(arrowstyle='-', color=colors['valid'], lw=0.5))
+
+    # 坐标轴
+    ax.set(xlabel='Epoch', ylabel='Loss', xlim=(0, len(train_loss) + 1))
+
+    # 样式
+    ax.legend(frameon=False, loc='upper right')
+    ax.grid(True, linestyle='--', alpha=0.3)
+    ax.spines[['top', 'right']].set_visible(False)
+
+    plt.tight_layout()
+    plt.savefig('./result/loss.png', dpi=300, bbox_inches='tight')
+    plt.close()
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    config_file_path = os.path.join(current_path, 'conf/config.yaml')
+    cfg = YParams(config_file_path, 'model')
+    cfg_data = YParams(config_file_path, "datapipe")
+
+    train_loss = np.load('./data/checkpoints/trloss.npy')
+    valid_loss = np.load('./data/checkpoints/valoss.npy')
+    plot_loss(train_loss, valid_loss)
+
+    data_dir = cfg_data.dataset.data_dir
+    total_files, channel_indices, time_step = get_metadata(data_dir, cfg_data.dataset.channels)
+
+    # Load data & Compute RMSE/ACC per channel
+    h5_files = sorted(glob.glob(os.path.join(data_dir, "data", "*.h5")))
+    with h5py.File(h5_files[0], "r") as f:
+        mu = f["global_means"][:]
+    clim_mean = mu[:, channel_indices, :, :]
+    get_result(total_files, channel_indices, time_step, data_dir, clim_mean)
+    show_result()
+
+    ##### 默认绘制 test_time 第一年的第一个时间步，用户可自行指定日期和变量 #####
+    test_year = cfg_data.dataset.test_time[0]
+    eg_files = [f'{test_year}010206']
+    channel_index = [cfg_data.dataset.channels.index(v) for v in ['2m_temperature', 'geopotential_500', 'temperature_500']]
+
+    selected_var = [cfg_data.dataset.channels[int(i)] for i in channel_index]
+    print(f"seleted date: {eg_files}")
+    print(f"selected channels: {selected_var}")
+    for file in eg_files:
+        year = file[:4]
+        t_idx = filename_to_index(file, time_step)
+        with h5py.File(os.path.join(data_dir, 'data', f'{year}.h5'), "r") as f:
+            label = f["fields"][t_idx]  # [C, H, W]
+            label = label[channel_indices]
+        pred = np.load(f'result/output/{file}.npy').squeeze()
+        for i in range(len(selected_var)):
+            filename = f'./result/{file}_{selected_var[i]}.png'
+            plot(label[channel_index[i]], pred[channel_index[i]], selected_var[i], filename)
+            print(f'✅plot {filename}')

scripts/train.py

@@ -0,0 +1,236 @@
+import torch
+import os
+import sys
+from pathlib import Path
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+import numpy as np
+import torch.distributed as dist
+import logging
+import time
+from tqdm import tqdm
+from torch.nn.parallel import DistributedDataParallel
+from model.fengwu import Fengwu
+from onescience.datapipes.climate import ERA5Datapipe
+from onescience.utils.YParams import YParams
+from onescience.memory.checkpoint import replace_function
+from onescience.utils.fcn.darcy_loss import LpLoss
+from apex import optimizers
+
+
+def loss_func(x, y):
+    return torch.nn.functional.l1_loss(x, y)
+
+
+def main():
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+    logger = logging.getLogger()
+
+    ## Model config init
+    config_file_path = os.path.join(current_path, "conf/config.yaml")
+    cfg = YParams(config_file_path, "model")
+
+    ## Distributed config init
+    cfg.world_size = 1
+    if "WORLD_SIZE" in os.environ:
+        cfg.world_size = int(os.environ["WORLD_SIZE"])
+    world_rank = 0
+    local_rank = 0
+    if cfg.world_size > 1:
+        dist.init_process_group(backend="nccl", init_method="env://")
+        local_rank = int(os.environ["LOCAL_RANK"])
+        world_rank = dist.get_rank()
+
+    ## DataLoader init
+    cfg_data = YParams(config_file_path, "datapipe")
+    datapipe = ERA5Datapipe(
+        dataset_dir=cfg_data.dataset.data_dir,
+        used_variables=cfg_data.dataset.channels,
+        used_years=cfg_data.dataset.train_time,
+        distributed=dist.is_initialized()
+    )
+    train_dataloader, train_sampler = datapipe.get_dataloader("train")
+    datapipe = ERA5Datapipe(
+        dataset_dir=cfg_data.dataset.data_dir,
+        used_variables=cfg_data.dataset.channels,
+        used_years=cfg_data.dataset.val_time,
+        distributed=dist.is_initialized()
+    )
+    val_dataloader, val_sampler = datapipe.get_dataloader("valid")
+
+    ## Model init
+    model = Fengwu(img_size=cfg_data.dataset.img_size,
+                   pressure_level=cfg.pressure_level,
+                   embed_dim=cfg.embed_dim,
+                   patch_size=cfg.patch_size,
+                   num_heads=cfg.num_heads,
+                   window_size=cfg.window_size,
+                   ).to(local_rank)
+    optimizer = optimizers.FusedAdam(model.parameters(), lr=cfg.lr)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,factor=0.2,patience=5,mode="min")
+    loss_obj = LpLoss()
+
+     ## Train process init
+    os.makedirs(cfg.checkpoint_dir, exist_ok=True)
+    train_loss_file = f"{cfg.checkpoint_dir}/trloss.npy"
+    valid_loss_file = f"{cfg.checkpoint_dir}/valoss.npy"
+    best_valid_loss = 1.0e6
+    best_loss_epoch = 0
+    train_losses = np.empty((0,), dtype=np.float32)
+    valid_losses = np.empty((0,), dtype=np.float32)
+
+    ## Get model params count
+    if cfg.world_size == 1:
+        total_params = sum(p.numel() for p in model.parameters())
+        print("\n\n")
+        print("-" * 50)
+        print(f"📂 now params is {total_params}, {total_params / 1e6:.2f}M, {total_params / 1e9:.2f}B")
+        print("-" * 50, "\n")
+
+    ## Load model weight if there exist well-trained model 
+    if os.path.exists(f"{cfg.checkpoint_dir}/model_bak.pth"):
+        if world_rank == 0:
+            print("\n\n")
+            print("-" * 50)
+            print(f"✅ There has a model weight, load and continue training...")
+            print(f'If you want to train a new model, ensure there is no *.pth file in {cfg.checkpoint_dir}')
+            print("-" * 50, "\n")
+        ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location=f'cuda:{local_rank}', weights_only=False)
+        model.load_state_dict(ckpt["model_state_dict"])
+        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+        scheduler.load_state_dict(ckpt["scheduler_state_dict"])
+        best_valid_loss = ckpt["best_valid_loss"]
+        best_loss_epoch = ckpt["best_loss_epoch"]
+        train_losses = np.load(train_loss_file)
+        valid_losses = np.load(valid_loss_file)
+
+    ## Distributed model
+    if cfg.world_size > 1:
+        model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)
+
+    world_rank == 0 and logger.info(f"start training ...")
+
+    for epoch in range(cfg.max_epoch):
+        if dist.is_initialized():
+            train_sampler.set_epoch(epoch)
+            val_sampler.set_epoch(epoch)
+
+        model.train()
+        train_loss = 0
+        start_time = time.time()
+        for j, data in enumerate(train_dataloader):
+            invar = data[0].to(local_rank, dtype=torch.float32)
+            outvar = data[1].to(local_rank, dtype=torch.float32)
+            surface = invar[:, :4, :, :]
+            z = invar[:, 4:41, :, :]
+            r = invar[:, 41:78, :, :]
+            u = invar[:, 78:115, :, :]
+            v = invar[:, 115:152, :, :]
+            t = invar[:, 152:189, :, :]
+
+            with replace_function(model,
+                        ["encoder_surface","encoder_z","encoder_r","encoder_u","encoder_v","encoder_t","fuser"],
+                        cfg.world_size > 1,):
+                surface_p, z_p, r_p, u_p, v_p, t_p = model(surface, z, r, u, v, t)
+
+            outvar_pred = torch.concat([surface_p, z_p, r_p, u_p, v_p, t_p],dim=1)
+
+            loss = loss_obj(outvar, outvar_pred)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            train_loss += loss.item()
+            if world_rank == 0:
+                logger.info(f'Train: Epoch {epoch}-{j+1}/{len(train_dataloader)} '
+                            f'[cost {int((time.time()-start_time) // 60):02}:{int((time.time()-start_time) % 60):02}] '
+                            f'[{(time.time()-start_time)/(j+1): .02f}s/{cfg_data.dataloader.batch_size}batch] '
+                            f'loss:{train_loss / (j+1): .04f}')
+            
+                          
+        train_loss /= len(train_dataloader)
+
+        model.eval()
+        valid_loss = 0
+        val_batch_time = time.time()
+        with torch.no_grad():
+            for j, data in enumerate(val_dataloader):
+                invar = data[0].to(local_rank, dtype=torch.float32)
+                outvar = data[1].to(local_rank, dtype=torch.float32)
+                surface = invar[:, :4, :, :]
+                z = invar[:, 4:41, :, :]
+                r = invar[:, 41:78, :, :]
+                u = invar[:, 78:115, :, :]
+                v = invar[:, 115:152, :, :]
+                t = invar[:, 152:189, :, :]
+
+                surface, z, r, u, v, t = model(surface, z, r, u, v, t)
+
+                outvar_pred = torch.concat(
+                    [surface_p, z_p, r_p, u_p, v_p, t_p], dim=1)
+
+                loss = loss_obj(outvar, outvar_pred)
+
+                if cfg.world_size > 1:
+                    loss_tensor = loss.detach().to(local_rank)
+                    dist.all_reduce(loss_tensor)
+                    loss = loss_tensor.item() / cfg.world_size
+                    valid_loss += loss
+                else:
+                    valid_loss += loss.item()
+
+                if world_rank == 0:
+                    logger.info(f'Valid: Epoch {epoch}-{j+1}/{len(val_dataloader)} '
+                            f'[cost {int((time.time()-start_time) // 60):02}:{int((time.time()-start_time) % 60):02}] '
+                            f'[{(time.time()-start_time)/(j+1): .02f}s/{cfg_data.dataloader.batch_size}batch] '
+                            f'loss:{valid_loss / (j+1): .04f}')
+    
+
+        valid_loss /= len(val_dataloader)
+        is_save_ckp = False
+        if valid_loss < best_valid_loss:
+            best_valid_loss = valid_loss
+            best_loss_epoch = epoch
+            world_rank == 0 and save_checkpoint(model, optimizer, scheduler, best_valid_loss, best_loss_epoch, cfg.checkpoint_dir)
+            is_save_ckp = True
+
+        scheduler.step(valid_loss)
+
+        if world_rank == 0:
+            logger.info(f"Epoch [{epoch + 1}/{cfg.max_epoch}], "
+                        f"Train Loss: {train_loss:.4f}, "
+                        f"Valid Loss: {valid_loss:.4f}, "
+                        f"Best loss at Epoch: {best_loss_epoch + 1}"
+                        + (", saving checkpoint" if is_save_ckp else "")
+                        )
+            train_losses = np.append(train_losses, train_loss)
+            valid_losses = np.append(valid_losses, valid_loss)
+
+            np.save(train_loss_file, train_losses)
+            np.save(valid_loss_file, valid_losses)
+
+        if epoch - best_loss_epoch > cfg.patience:
+            print(f"Loss has not decrease in {cfg.patience} epochs, stopping training...")
+            exit()
+
+
+def save_checkpoint(model, optimizer, scheduler, best_valid_loss,
+                    best_loss_epoch, model_path):
+    model_to_save = model.module if hasattr(model, "module") else model
+    state = {
+        "model_state_dict": model_to_save.state_dict(),
+        "optimizer_state_dict": optimizer.state_dict(),
+        "scheduler_state_dict": scheduler.state_dict(),
+        "best_valid_loss": best_valid_loss,
+        "best_loss_epoch": best_loss_epoch,
+    }
+    torch.save(state, f"{model_path}/model.pth")
+    ### the weight file saving may interrupted due to DCU queue limit, get a backup to ensure there at least has one model 
+    os.system(f"mv {model_path}/model.pth {model_path}/model_bak.pth")
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    main()