Upload folder using huggingface_hub

README.md

@@ -0,0 +1,140 @@
+---
+license: apache-2.0
+tasks:
+  - weather-forecast
+frameworks:
+  - pytorch
+language:
+  - en
+  - zh
+tags:
+  - OneScience
+  - Earth Science
+  - Weather Forecast
+  - Short-to-Medium-Range Weather Forecast
+  - ERA5
+datasets:
+ - OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Pangu-Weather</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+Pangu-Weather is a global medium-range weather forecast model proposed by Huawei Cloud, capable of rapidly predicting surface variables and multi-pressure-level upper-air variables.
+
+Paper: Accurate medium-range global weather forecasting with 3D neural networks  
+https://www.nature.com/articles/s41586-023-06185-3
+
+# Model Description
+
+Pangu-Weather is based on a 3D Earth-Specific Transformer architecture, trained on ERA5 data, and designed for short-to-medium-range weather forecasting.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Weather Forecast Training | Train Pangu-Weather 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/Pangu-Weather --local_dir ./pangu_weather
+cd pangu_weather
+```
+
+### 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 ERA5 data from 1979 to 2025 in the `weights/` folder. At the 6-hour forecast lead time, these weights outperform the official open-source ONNX weights.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference results will be saved to `result/output/`.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+# 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 Pangu-Weather paper.

conf/config.yaml

@@ -0,0 +1,130 @@
+model:
+    
+    start_epoch: 0
+    max_epoch: 1
+    patch_size: [2, 4, 4]
+    embed_dim: 192
+    num_heads: [6, 12, 12, 6]
+    window_size: [2, 6, 12]
+    patience: 50
+    checkpoint_dir: "./data/checkpoints"
+
+datapipe:
+  name: "ERA5"
+  task: "weather_forecasting"
+
+  # dataset设定
+  dataset:
+    type: "hdf5"
+    data_dir: './data/'
+    
+    train_time: [1979]
+    val_time: [2021]
+    test_time: [2023]
+
+    img_size: [721, 1440]
+    verbose: true
+    cache: false
+    # 气象变量
+    channels:  ['mean_sea_level_pressure', '10m_u_component_of_wind', '10m_v_component_of_wind', '2m_temperature', 
+                
+                'geopotential_1000', 'geopotential_925', 'geopotential_850', 'geopotential_700', 'geopotential_600', 'geopotential_500', 
+                'geopotential_400', 'geopotential_300', 'geopotential_250', 'geopotential_200', 'geopotential_150', 'geopotential_100', 'geopotential_50',
+                
+                'specific_humidity_1000', 'specific_humidity_925', 'specific_humidity_850', 'specific_humidity_700', 'specific_humidity_600',
+                'specific_humidity_500', 'specific_humidity_400', 'specific_humidity_300', 'specific_humidity_250', 'specific_humidity_200', 
+                'specific_humidity_150', 'specific_humidity_100', 'specific_humidity_50', 
+                
+                'temperature_1000', 'temperature_925', 'temperature_850', 'temperature_700', 'temperature_600', 'temperature_500', 'temperature_400', 
+                'temperature_300', 'temperature_250', 'temperature_200','temperature_150', 'temperature_100', 'temperature_50', 
+                
+                'u_component_of_wind_1000', 'u_component_of_wind_925', 'u_component_of_wind_850', 'u_component_of_wind_700', 'u_component_of_wind_600', 
+                'u_component_of_wind_500', 'u_component_of_wind_400', 'u_component_of_wind_300', 'u_component_of_wind_250', 'u_component_of_wind_200', 
+                'u_component_of_wind_150', 'u_component_of_wind_100', 'u_component_of_wind_50', 
+                
+                'v_component_of_wind_1000', 'v_component_of_wind_925', 'v_component_of_wind_850', 'v_component_of_wind_700', 'v_component_of_wind_600', 
+                'v_component_of_wind_500', 'v_component_of_wind_400', 'v_component_of_wind_300', 'v_component_of_wind_250', 'v_component_of_wind_200', 
+                'v_component_of_wind_150', 'v_component_of_wind_100', 'v_component_of_wind_50'
+                ]
+    weights: [  # ---- 4 surface variables ----
+                1.50,  # mean_sea_level_pressure
+                0.77,  # 10m_u_component_of_wind
+                0.66,  # 10m_v_component_of_wind
+                3.00,  # 2m_temperature
+                # ---- 13 geopotential_* (Z) ----
+                3.00, 3.00, 3.00, 3.00, 3.00, 3.00, 3.00, 3.00, 3.00, 3.00, 3.00, 3.00, 3.00,
+                # ---- 13 specific_humidity_* (Q) ----
+                0.60, 0.60, 0.60, 0.60, 0.60, 0.60, 0.60, 0.60, 0.60, 0.60, 0.60, 0.60, 0.60,
+                # ---- 13 temperature_* (T) ----
+                1.50, 1.50, 1.50, 1.50, 1.50, 1.50, 1.50, 1.50, 1.50, 1.50, 1.50, 1.50, 1.50,
+                # ---- 13 u_component_of_wind_* (U) ----
+                0.77, 0.77, 0.77, 0.77, 0.77, 0.77, 0.77, 0.77, 0.77, 0.77, 0.77, 0.77, 0.77,
+                # ---- 13 v_component_of_wind_* (V) ----
+                0.54, 0.54, 0.54, 0.54, 0.54, 0.54, 0.54, 0.54, 0.54, 0.54, 0.54, 0.54, 0.54,
+            ]
+    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/pangu.py

@@ -0,0 +1,218 @@
+import math
+import torch
+import numpy as np
+
+from torch import nn
+from dataclasses import dataclass
+from onescience.models.meta import ModelMetaData
+
+from onescience.modules.embedding.panguembedding import PanguEmbedding
+from onescience.modules.fuser.pangufuser import PanguFuser
+from onescience.modules.recovery.pangupatchrecovery import PanguPatchRecovery
+from onescience.modules.sample.pangudownsample import PanguDownSample
+from onescience.modules.sample.panguupsample import PanguUpSample
+
+@dataclass
+class MetaData(ModelMetaData):
+    name: str = "Pangu"
+    # 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 Pangu(nn.Module):
+    """
+    Pangu-Weather 的主模型实现。
+
+    该模型使用以下组件完成编码、主干特征提取与输出恢复：
+
+    - `OneEmbedding(style="PanguEmbedding")`
+      - 将 surface 分支与 upper-air 分支分别映射到 patch 特征空间
+    - `OneFuser(style="PanguFuser")`
+      - 在统一的三维 token 网格上完成主干特征融合
+    - `OneSample(style="PanguDownSample" / "PanguUpSample")`
+      - 在主干中完成空间尺度变换
+    - `OneRecovery(style="PanguPatchRecovery")`
+      - 将 patch 级特征恢复为目标物理场
+
+    在当前实现中：
+
+    - 输入包含 7 个 surface 通道，其中 4 个是待预测地表变量，3 个是静态掩码
+    - 输入包含 5 个 upper-air 变量，每个变量对应 13 个气压层
+    - surface patch token 会先在 `PressureLevels` 维补成长度为 1
+    - 然后与 upper-air patch token 沿 `PressureLevels` 维拼接
+    - 之后统一送入三维 `PanguFuser` 主干
+
+    Reference:
+    - `Pangu-Weather: A 3D High-Resolution Model for Fast and Accurate Global Weather Forecast`
+    - https://arxiv.org/abs/2211.02556
+
+    Args:
+        img_size (tuple[int, int]):
+            输入空间尺寸 `(Height, Width)`。
+        patch_size (tuple[int, int, int]):
+            patch 切分尺寸 `(PatchPressureLevels, PatchHeight, PatchWidth)`。
+        embed_dim (int):
+            patch embedding 后的特征维度。
+        num_heads (tuple[int, int, int, int]):
+            四个主干阶段对应的注意力头数。
+        window_size (tuple[int, int, int]):
+            三维窗口大小 `(PressureLevelsWindow, HeightWindow, WidthWindow)`。
+    """
+
+    def __init__(
+        self,
+        img_size=(721, 1440),
+        patch_size=(2, 4, 4),
+        embed_dim=192,
+        num_heads=(6, 12, 12, 6),
+        window_size=(2, 6, 12),
+    ):
+        super().__init__()
+        drop_path = np.linspace(0, 0.2, 8).tolist()
+        # Surface input contains 4 predicted variables and 3 static masks:
+        # topography mask, land-sea mask, and soil type mask.
+
+        self.patchembed2d = PanguEmbedding(
+            img_size=img_size,
+            patch_size=patch_size[1:],
+            Variables=7,
+            embed_dim=embed_dim,
+        )
+        self.patchembed3d = PanguEmbedding(
+            img_size=(13, *img_size),
+            patch_size=patch_size,
+            Variables=5,
+            embed_dim=embed_dim,
+        )
+
+        patched_input_shape = (
+            8,
+            math.ceil(img_size[0] / patch_size[1]),
+            math.ceil(img_size[1] / patch_size[2]),
+        )
+
+        self.layer1 = PanguFuser(
+            dim=embed_dim,
+            input_resolution=patched_input_shape,
+            depth=2,
+            num_heads=num_heads[0],
+            window_size=window_size,
+            drop_path=drop_path[:2],
+        )
+
+        patched_downsampled_shape = (
+            8,
+            math.ceil(patched_input_shape[1] / 2),
+            math.ceil(patched_input_shape[2] / 2),
+        )
+
+        self.downsample = PanguDownSample(
+            in_dim=embed_dim,
+            input_resolution=patched_input_shape,
+            output_resolution=patched_downsampled_shape,
+        )
+        self.layer2 = PanguFuser(
+            dim=embed_dim * 2,
+            input_resolution=patched_downsampled_shape,
+            depth=6,
+            num_heads=num_heads[1],
+            window_size=window_size,
+            drop_path=drop_path[2:],
+        )
+        self.layer3 = PanguFuser(
+            dim=embed_dim * 2,
+            input_resolution=patched_downsampled_shape,
+            depth=6,
+            num_heads=num_heads[2],
+            window_size=window_size,
+            drop_path=drop_path[2:],
+        )
+        self.upsample = PanguUpSample(
+            in_dim=embed_dim * 2,
+            out_dim=embed_dim,
+            input_resolution=patched_downsampled_shape,
+            output_resolution=patched_input_shape,
+        )
+        self.layer4 = PanguFuser(
+            dim=embed_dim,
+            input_resolution=patched_input_shape,
+            depth=2,
+            num_heads=num_heads[3],
+            window_size=window_size,
+            drop_path=drop_path[:2],
+        )
+
+        # The recovered surface output contains only the 4 prognostic surface variables.
+        # Static masks are input-only features and are not part of the prediction target.
+        self.patchrecovery2d = PanguPatchRecovery(
+            img_size=(721, 1440),
+            patch_size=(4, 4),
+            in_chans=embed_dim * 2,
+            out_chans=4,
+        )
+        self.patchrecovery3d = PanguPatchRecovery(
+            img_size=(13, 721, 1440),
+            patch_size=(2, 4, 4),
+            in_chans=embed_dim * 2,
+            out_chans=5,
+        )
+
+    def forward(self, x):
+        """
+        Args:
+            x (torch.Tensor):
+                Input tensor with shape `(Batch, 4 + 3 + 5 * 13, Height, Width)`.
+
+                Channel layout:
+                - first 4 channels: prognostic surface variables
+                - next 3 channels: static masks
+                - remaining `5 * 13` channels: upper-air variables flattened over pressure levels
+
+        Returns:
+            tuple[torch.Tensor, torch.Tensor]:
+                - surface output:
+                  `(Batch, 4, Height, Width)`
+                - upper-air output:
+                  `(Batch, 5, 13, Height, Width)`
+        """
+        SurfaceInput = x[:, :7, :, :]
+        UpperAirInput = x[:, 7:, :, :].reshape(x.shape[0], 5, 13, x.shape[2], x.shape[3])
+
+        SurfaceFeatures = self.patchembed2d(SurfaceInput)
+        UpperAirFeatures = self.patchembed3d(UpperAirInput)
+
+        CombinedFeatures = torch.concat(
+            [SurfaceFeatures.unsqueeze(2), UpperAirFeatures], dim=2
+        )
+        Batch, Channels, PressureLevels, Height, Width = CombinedFeatures.shape
+        Tokens = CombinedFeatures.reshape(Batch, Channels, -1).transpose(1, 2)
+
+        Tokens = self.layer1(Tokens)
+        SkipTokens = Tokens
+
+        Tokens = self.downsample(Tokens)
+        Tokens = self.layer2(Tokens)
+        Tokens = self.layer3(Tokens)
+        Tokens = self.upsample(Tokens)
+        Tokens = self.layer4(Tokens)
+
+        OutputFeatures = torch.concat([Tokens, SkipTokens], dim=-1)
+        OutputFeatures = OutputFeatures.transpose(1, 2).reshape(
+            Batch, -1, PressureLevels, Height, Width
+        )
+        output_surface = OutputFeatures[:, :, 0, :, :]
+        output_upper_air = OutputFeatures[:, :, 1:, :, :]
+
+        output_surface = self.patchrecovery2d(output_surface)
+        output_upper_air = self.patchrecovery3d(output_upper_air)
+        return output_surface, output_upper_air

scripts/fake_data.py

@@ -0,0 +1,95 @@
+import os
+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,90 @@
+import sys
+from pathlib import Path
+
+# 获取项目根目录（train.py上级的上级）
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+import torch
+import os
+import sys
+import glob
+import numpy as np
+import h5py
+from tqdm import tqdm
+from model.pangu import Pangu
+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")
+
+    static_dir = os.path.join(cfg_data.dataset.data_dir, "static")
+    
+    land_mask = torch.from_numpy(np.load(os.path.join(static_dir, "land_mask.npy")).astype(np.float32))
+    soil_type = torch.from_numpy(np.load(os.path.join(static_dir, "soil_type.npy")).astype(np.float32))
+    topography = torch.from_numpy(np.load(os.path.join(static_dir, "topography.npy")).astype(np.float32))
+    topography = (topography - topography.mean()) / (topography.std(unbiased=False) + 1e-6)
+    surface_mask = torch.stack([land_mask, soil_type, topography], dim=0).to('cuda:0')
+    surface_mask = surface_mask.unsqueeze(0).repeat(cfg_data.dataloader.batch_size, 1, 1, 1)
+
+    ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location="cuda:0")
+    model = Pangu(img_size=cfg_data.dataset.img_size,
+                  patch_size=cfg.patch_size,
+                  embed_dim=cfg.embed_dim,
+                  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]
+            outvar = data[1]
+            filename = data[4][-1][0]
+            invar_surface = invar[:, :4, :, :].to("cuda:0", dtype=torch.float32)
+            invar_upper_air = invar[:, 4:, :, :].to("cuda:0", dtype=torch.float32)
+            invar = torch.concat([invar_surface, surface_mask, invar_upper_air], dim=1)
+
+            out_surface, out_upper_air = model(invar)
+            out_upper_air = out_upper_air.reshape(invar_upper_air.shape)
+            pred_var = torch.concat([out_surface, out_upper_air], 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,226 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import os
+import sys
+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 group_files_by_year(total_files, time_step):
+    """将输出文件按年份归组，减少重复打开年度 h5 文件的开销"""
+    files_by_year = {}
+    for file in total_files:
+        fname = file[:-4]
+        year = fname[:4]
+        files_by_year.setdefault(year, []).append((file, filename_to_index(fname, time_step)))
+    return files_by_year
+
+
+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))
+        files_by_year = group_files_by_year(total_files, time_step)
+        with tqdm(total=len(total_files), unit="files") as pbar:
+            for year, year_files in files_by_year.items():
+                with h5py.File(os.path.join(data_dir, 'data', f'{year}.h5'), "r") as f:
+                    fields = f["fields"]
+                    for file, t_idx in year_files:
+                        label = 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)))
+                        pbar.update(1)
+        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,248 @@
+import sys
+from pathlib import Path
+
+# 获取项目根目录（train.py上级的上级）
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+
+import torch
+import os
+import numpy as np
+import torch.distributed as dist
+import logging
+import time
+import torch.nn.functional as F
+from torch.nn.parallel import DistributedDataParallel
+from model.pangu import Pangu
+from onescience.datapipes.climate import ERA5Datapipe
+from onescience.utils.YParams import YParams
+from onescience.memory.checkpoint import replace_function
+from apex import optimizers
+
+
+
+
+def loss_func(x, y, weights, level_weight=1.0):
+    return level_weight * (F.l1_loss(x, y, reduction='none') * weights).mean()
+
+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(),
+        batch_size=cfg_data.dataloader.batch_size,
+        num_workers=cfg_data.dataloader.num_workers
+    )
+    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(),
+        batch_size=cfg_data.dataloader.batch_size,
+        num_workers=cfg_data.dataloader.num_workers
+    )
+    val_dataloader, val_sampler = datapipe.get_dataloader("valid")
+
+    surface_weights = torch.as_tensor(cfg_data.dataset.weights[:4], device=local_rank, dtype=torch.float32).view(1, -1, 1, 1)
+    pressure_weights = torch.as_tensor(cfg_data.dataset.weights[4:], device=local_rank, dtype=torch.float32).view(1, -1, 1, 1)
+
+    static_dir = os.path.join(cfg_data.dataset.data_dir, "static")
+    
+    land_mask = torch.from_numpy(np.load(os.path.join(static_dir, "land_mask.npy")).astype(np.float32))
+    soil_type = torch.from_numpy(np.load(os.path.join(static_dir, "soil_type.npy")).astype(np.float32))
+    topography = torch.from_numpy(np.load(os.path.join(static_dir, "topography.npy")).astype(np.float32))
+    topography = (topography - topography.mean()) / (topography.std(unbiased=False) + 1e-6)
+    surface_mask = torch.stack([land_mask, soil_type, topography], dim=0).to(local_rank)
+    surface_mask = surface_mask.unsqueeze(0).repeat(cfg_data.dataloader.batch_size, 1, 1, 1)
+
+    ## Model init
+    model = Pangu(img_size=cfg_data.dataset.img_size,
+                  patch_size=cfg.patch_size,
+                  embed_dim=cfg.embed_dim,
+                  num_heads=cfg.num_heads,
+                  window_size=cfg.window_size,
+                  ).to(local_rank)
+    optimizer = optimizers.FusedAdam(model.parameters(), betas=(0.9, 0.999), lr=5e-4, weight_decay=3e-6)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
+    
+    ## 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)
+    current_epoch = 0
+
+    ## 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"]
+        current_epoch = ckpt["current_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)
+
+    world_rank == 0 and logger.info(f"start training ...")
+
+    for epoch in range(current_epoch, 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]
+            outvar = data[1]
+            invar_surface = invar[:, :4, :, :].to(local_rank, dtype=torch.float32)
+            invar_upper_air = invar[:, 4:, :, :].to(local_rank, dtype=torch.float32)
+            invar = torch.concat([invar_surface, surface_mask, invar_upper_air], dim=1)
+            tar_surface = outvar[:, :4, :, :].to(local_rank, dtype=torch.float32)
+            tar_upper_air = outvar[:, 4:, :, :].to(local_rank, dtype=torch.float32)
+
+            with replace_function(model,["layer2", "layer3"],cfg.world_size > 1):
+                out_surface, out_upper_air = model(invar)
+
+            out_upper_air = out_upper_air.reshape(tar_upper_air.shape)
+            loss1 = loss_func(out_surface, tar_surface, surface_weights,  level_weight=0.25)
+            loss2 = loss_func(out_upper_air, tar_upper_air, pressure_weights, level_weight=1.0)
+            # 总 loss
+            loss = loss1 + loss2
+            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
+        with torch.no_grad():
+            start_time = time.time()
+            for j, data in enumerate(val_dataloader):
+                invar = data[0]
+                outvar = data[1]
+                invar_surface = invar[:, :4, :, :].to(local_rank, dtype=torch.float32)
+                invar_upper_air = invar[:, 4:, :, :].to(local_rank, dtype=torch.float32)
+                invar = torch.concat([invar_surface, surface_mask, invar_upper_air], dim=1)
+                tar_surface = outvar[:, :4, :, :].to(local_rank, dtype=torch.float32)
+                tar_upper_air = outvar[:, 4:, :, :].to(local_rank, dtype=torch.float32)
+
+                with replace_function(model,["layer2", "layer3"],cfg.world_size > 1):
+                    out_surface, out_upper_air = model(invar)
+
+                out_upper_air = out_upper_air.reshape(tar_upper_air.shape)
+                loss1 = loss_func(out_surface, tar_surface, surface_weights,  level_weight=0.25).item()
+                loss2 = loss_func(out_upper_air, tar_upper_air, pressure_weights, level_weight=1.0).item()
+                # 总 loss
+                loss = loss1 + loss2
+
+                if cfg.world_size > 1:
+                    loss_tensor = torch.tensor(loss, device=local_rank)
+                    dist.all_reduce(loss_tensor)
+                    loss = loss_tensor.item() / cfg.world_size
+                valid_loss += loss
+                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, epoch)
+            is_save_ckp = True
+
+        scheduler.step()
+
+        if world_rank == 0:
+            logger.info(f"Epoch [{epoch}/{cfg.max_epoch}], "
+                        f"Train Loss: {train_loss:.4f}, "
+                        f"Valid Loss: {valid_loss:.4f}, "
+                        f"Best loss at Epoch: {best_loss_epoch}"
+                        + (", 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, epoch):
+    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,
+             "current_epoch": 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()

weight/.gitkeep

@@ -0,0 +1 @@
+