Upload folder using huggingface_hub

README.md

@@ -0,0 +1,173 @@
+---
+frameworks: PyTorch
+language:
+- en
+- zh
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Weather Forecast
+- Short-to-Medium-Range Weather Forecast
+- ERA5
+tasks: []
+datasets:
+  - OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">FuXi</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+FuXi is a global weather forecast foundation model jointly developed by Fudan University and multiple institutions. It is the first end-to-end machine learning framework capable of independently performing data assimilation (DA) and cyclic forecasting.
+
+Paper: FuXi: A cascade machine learning forecasting system for 15-day global weather forecast
+
+https://arxiv.org/abs/2306.12873
+
+# Model Description
+
+The FuXi model is trained through a three-stage cascaded approach: short → medium → long. Its training input primarily consists of ERA5 reanalysis data.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Weather Forecast Training | Train FuXi (short/medium/long three stages) 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/FuXi --local_dir ./FuXi
+cd FuXi
+```
+
+### 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
+
+FuXi consists of 3 stages and **must be executed in order**. The inference result of each stage serves as the input for the next stage:
+
+**short (train) → short (inference) → medium (train) → medium (inference) → long (train) → long (inference)**
+
+**1) Train the short model (train from scratch, as the starting entry point)**
+
+Single GPU:
+
+```bash
+python scripts/train_short.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_short.py
+```
+
+**2) Short inference (generate input data for medium)**
+
+```bash
+python scripts/inference.py short
+```
+
+**3) Train the medium model (requires short weights + short inference results)**
+
+```bash
+python scripts/train_medium.py
+```
+
+**4) Medium inference (generate input data for long)**
+
+```bash
+python scripts/inference.py medium
+```
+
+**5) Train the long model (requires medium weights + medium inference results)**
+
+```bash
+python scripts/train_long.py
+```
+
+### 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
+
+Each stage can perform inference independently:
+
+```bash
+python scripts/inference.py short
+python scripts/inference.py medium
+python scripts/inference.py long
+```
+
+Inference results will be saved to `result/output/<stage>/`.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py short
+python scripts/result.py medium
+python scripts/result.py long
+```
+
+# 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 FuXi paper.

conf/config.yaml

@@ -0,0 +1,118 @@
+model:
+
+    start_epoch: 0
+    max_epoch: 1
+    train_lr: 2.5E-4
+    finetune_lr: 1E-7
+    patch_size: [2, 4, 4]
+    embed_dim: 1536
+    num_groups: 32
+    num_heads: 8 
+    window_size: 7
+    short_num_steps: 20
+    medium_num_steps: 40
+    long_num_steps: 60
+    finetune_step: 1
+    step_change_freq: 200
+    input_steps: 2
+    num_blocks: 8
+    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: [2022]
+    test_time: [2023] 
+    img_size: [2, 721, 1440]
+    verbose: true
+    cache: false
+    # 气象变量
+
+    channels:  ['10m_u_component_of_wind', '10m_v_component_of_wind', '2m_temperature', 'mean_sea_level_pressure', 'total_precipitation',
+                'geopotential_50', 'geopotential_100', 'geopotential_150', 'geopotential_200', 'geopotential_250', 
+                'geopotential_300', 'geopotential_400', 'geopotential_500', 'geopotential_600', 'geopotential_700', 
+                'geopotential_850', 'geopotential_925', 'geopotential_1000',
+                'relative_humidity_50',  'relative_humidity_100', 'relative_humidity_150', 'relative_humidity_200', 'relative_humidity_250', 
+                'relative_humidity_300', 'relative_humidity_400', 'relative_humidity_500', 'relative_humidity_600', 'relative_humidity_700', 
+                'relative_humidity_850','relative_humidity_925', 'relative_humidity_1000',
+                'u_component_of_wind_50', 'u_component_of_wind_100', 'u_component_of_wind_150', 'u_component_of_wind_200', 'u_component_of_wind_250', 
+                'u_component_of_wind_300', 'u_component_of_wind_400', 'u_component_of_wind_500', 'u_component_of_wind_600','u_component_of_wind_700', 
+                'u_component_of_wind_850', 'u_component_of_wind_925', 'u_component_of_wind_1000',
+                'v_component_of_wind_50', 'v_component_of_wind_100', 'v_component_of_wind_150', 'v_component_of_wind_200', 'v_component_of_wind_250', 
+                'v_component_of_wind_300', 'v_component_of_wind_400', 'v_component_of_wind_500', 'v_component_of_wind_600','v_component_of_wind_700', 
+                'v_component_of_wind_850', 'v_component_of_wind_925', 'v_component_of_wind_1000',
+                'temperature_50', 'temperature_100', 'temperature_150', 'temperature_200', 'temperature_250', 
+                'temperature_300', 'temperature_400', 'temperature_500', 'temperature_600', 'temperature_700', 
+                'temperature_850', 'temperature_925', '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/fuxi.py

@@ -0,0 +1,142 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from onescience.modules.embedding.fuxiembedding import FuxiEmbedding
+from onescience.modules.fc.fuxifc import FuxiFC
+from onescience.modules.transformer.fuxitransformer import FuxiTransformer
+
+
+class Fuxi(nn.Module):
+    """
+    Fuxi 的主模型实现。
+
+    该模型使用以下组件完成输入编码、二维 trunk 特征提取与 patch 级输出恢复：
+
+    - `OneEmbedding(style="FuxiEmbedding")`
+      - 将 `(TimeSteps, Height, Width)` 三维时空块映射为 patch 特征
+    - `OneTransformer(style="FuxiTransformer")`
+      - 在二维特征图上执行下采样、Swin trunk、上采样
+    - `OneFC(style="FuxiFC")`
+      - 将每个二维网格位置的 embedding 特征映射为 patch 级输出变量
+
+    在当前实现中：
+
+    - 输入包含多个时间步的二维气象场
+    - `patch_size[0]` 默认与 `TimeSteps` 相同，使 embedding 后时间维压缩为 1
+    - trunk 只处理二维特征图
+    - 最终通过 patch 重排与双线性插值恢复到目标空间分辨率
+
+    Args:
+        img_size (tuple[int, int, int]):
+            输入空间尺寸 `(TimeSteps, Height, Width)`。
+        patch_size (tuple[int, int, int]):
+            patch 切分尺寸 `(PatchTimeSteps, PatchHeight, PatchWidth)`。
+        in_chans (int):
+            输入变量通道数。
+        out_chans (int):
+            输出变量通道数。
+        embed_dim (int):
+            embedding 特征维度。
+        num_groups (int):
+            trunk 中采样模块的 `GroupNorm` 分组数。
+        num_heads (int):
+            `SwinTransformerV2Stage` 的注意力头数。
+        window_size (int | tuple[int, int]):
+            trunk 局部窗口大小。
+    """
+
+    def __init__(
+        self,
+        img_size=(2, 721, 1440),
+        patch_size=(2, 4, 4),
+        in_chans=70,
+        out_chans=70,
+        embed_dim=1536,
+        num_groups=32,
+        num_heads=8,
+        window_size=7,
+    ):
+        super().__init__()
+
+        TimeSteps, Height, Width = img_size
+        PatchTimeSteps, PatchHeight, PatchWidth = patch_size
+        if TimeSteps != PatchTimeSteps:
+            raise ValueError(
+                "Current Fuxi model expects patch_size[0] to equal img_size[0] "
+                "so the embedding output time dimension is 1 before squeeze"
+            )
+
+        EmbeddedHeight = Height // PatchHeight
+        EmbeddedWidth = Width // PatchWidth
+        TransformerInputResolution = (
+            EmbeddedHeight // 2,
+            EmbeddedWidth // 2,
+        )
+
+        self.cube_embedding = FuxiEmbedding(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
+        self.u_transformer = FuxiTransformer(
+            embed_dim=embed_dim,
+            num_groups=num_groups,
+            input_resolution=TransformerInputResolution,
+            num_heads=num_heads,
+            window_size=window_size,
+        )
+        self.fc = FuxiFC(
+            in_channels=embed_dim,
+            out_channels=out_chans * PatchHeight * PatchWidth,
+        )
+
+        self.patch_size = patch_size
+        self.transformer_input_resolution = TransformerInputResolution
+        self.embedded_resolution = (EmbeddedHeight, EmbeddedWidth)
+        self.out_chans = out_chans
+        self.img_size = img_size
+
+    def forward(self, x):
+        """
+        Args:
+            x (torch.Tensor):
+                输入张量，形状为 `(Batch, in_chans, TimeSteps, Height, Width)`。
+
+        Returns:
+            torch.Tensor:
+                输出张量，形状为 `(Batch, out_chans, Height, Width)`。
+        """
+        Batch, _, _, _, _ = x.shape
+        _, PatchHeight, PatchWidth = self.patch_size
+        EmbeddedHeight, EmbeddedWidth = self.embedded_resolution
+
+        x = self.cube_embedding(x)
+        if x.shape[2] != 1:
+            raise ValueError(
+                f"Expected embedding time dimension 1 before squeeze, but received {x.shape[2]}"
+            )
+
+        x = x.squeeze(2)
+        x = self.u_transformer(x)
+        x = self.fc(x.permute(0, 2, 3, 1))
+        x = x.reshape(
+            Batch,
+            EmbeddedHeight,
+            EmbeddedWidth,
+            PatchHeight,
+            PatchWidth,
+            self.out_chans,
+        ).permute(0, 1, 3, 2, 4, 5)
+        x = x.reshape(
+            Batch,
+            EmbeddedHeight * PatchHeight,
+            EmbeddedWidth * PatchWidth,
+            self.out_chans,
+        )
+        x = x.permute(0, 3, 1, 2)
+
+        x = F.interpolate(x, size=self.img_size[1:], mode="bilinear")
+
+        return x

scripts/data_loader.py

@@ -0,0 +1,205 @@
+import os
+import glob
+import h5py
+import pytz
+import numpy as np
+import torch
+
+from datetime import datetime, timedelta
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from onescience.datapipes.climate.utils.invariant import latlon_grid
+from onescience.datapipes.climate.utils.zenith_angle import cos_zenith_angle
+
+
+class ERA5Datapipe:
+    def __init__(
+        self,
+        dataset_dir,
+        used_years,
+        used_variables,
+        pattern='medium',
+        distributed=False,
+        input_steps=1,
+        output_steps=1,
+        normalize=True,
+        batch_size=1,
+        num_workers=4,
+    ):
+        self.dataset_dir   = dataset_dir
+        self.used_years    = used_years
+        self.used_variables = used_variables
+        self.pattern       = pattern
+        self.distributed   = distributed
+        self.input_steps   = input_steps
+        self.output_steps  = output_steps
+        self.normalize     = normalize
+        self.batch_size    = batch_size
+        self.num_workers   = num_workers
+
+    def get_dataloader(self, mode):
+        dataset = ERA5Dataset(
+            dataset_dir=self.dataset_dir,
+            used_years=self.used_years,
+            used_variables=self.used_variables,
+            pattern=self.pattern,
+            input_steps=self.input_steps,
+            output_steps=self.output_steps,
+            normalize=self.normalize,
+        )
+        is_train = (mode == 'train')
+
+        sampler = DistributedSampler(dataset, shuffle=is_train) if self.distributed else None
+
+        return DataLoader(
+            dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            shuffle=(is_train and not self.distributed),
+            sampler=sampler,
+            drop_last=self.distributed,
+        ), sampler
+
+
+class ERA5Dataset(Dataset):
+    def __init__(
+        self,
+        dataset_dir,
+        used_years,
+        used_variables,
+        pattern='medium',
+        input_steps=1,
+        output_steps=1,
+        normalize=True,
+    ):
+        self.dataset_dir    = dataset_dir
+        self.used_years     = used_years
+        self.used_variables = used_variables
+        self.pattern        = pattern
+        self.input_steps    = input_steps
+        self.output_steps   = output_steps
+        self.normalize      = normalize
+
+        self._init_avail_samples()
+        self._init_normalized_files()
+        self._init_npy_files()
+        self._init_latlon_grid()
+
+    def _init_avail_samples(self):
+        h5_files = sorted(glob.glob(os.path.join(self.dataset_dir, "data", "*.h5")))
+        available_years = [int(os.path.basename(f).replace(".h5", "")) for f in h5_files]
+
+        missing_years = [y for y in self.used_years if y not in available_years]
+        if missing_years:
+            raise ValueError(f"❌ Years not found in dataset: {missing_years}")
+
+        # ── 读取变量信息 & 校验 ───────────────────────────────
+        with h5py.File(h5_files[0], "r") as f:
+            ds = f["fields"]
+            self.T, self.C, self.H, self.W = ds.shape
+            all_variables = [v.decode() if isinstance(v, bytes) else v for v in ds.attrs["variables"]]
+            self.time_step = int(ds.attrs["time_step"])
+
+        missing_vars = [v for v in self.used_variables if v not in all_variables]
+        if missing_vars:
+            raise ValueError(f"❌ Variables not found in dataset: {missing_vars}")
+
+        # ── 建立索引 ──────────────────────────────────────────
+        self.channel_indices = [all_variables.index(v) for v in self.used_variables]
+
+    def _init_normalized_files(self):
+        # ── 从每年 h5 内嵌的 global_means/global_stds 读取归一化统计量 ──
+        h5_files = sorted(glob.glob(os.path.join(self.dataset_dir, "data", "*.h5")))
+        with h5py.File(h5_files[0], "r") as f:
+            mu  = f["global_means"][:]   # [1, C, 1, 1]
+            std = f["global_stds"][:]
+        self.mu = torch.as_tensor(mu[:, self.channel_indices, :, :], dtype=torch.float32)
+        self.sd = torch.as_tensor(std[:, self.channel_indices, :, :], dtype=torch.float32)
+
+    def _init_npy_files(self):
+        """读取前一阶段模型输出的 npy 文件列表（作为 invar 输入）"""
+        self.files = {}
+        for year in self.used_years:
+            if self.pattern == 'medium':
+                path = os.path.join('./result/short/data/', str(year))
+            else:
+                path = os.path.join('./result/medium/data', str(year))
+            files = sorted(glob.glob(os.path.join(path, "*.npy")))
+            if not files:
+                raise ValueError(f"❌ No npy files found for year {year} under {path}")
+            self.files[year] = files
+
+        n_files = len(files)
+        self.samples_per_year = n_files - self.output_steps - (self.input_steps - 1)
+        self.total_samples = len(self.used_years) * self.samples_per_year
+
+        if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
+            print('\n')
+            print('-' * 50)
+            print(f"📂 Pattern: {self.pattern}, used: {self.used_years} years")
+            print(f'📂 each year contains {self.samples_per_year} usable samples '
+                  f'({n_files} files, input {self.input_steps}, output {self.output_steps})')
+            print(f'📂 total usable samples: {self.total_samples}')
+            print('-' * 50, '\n')
+
+    def _init_latlon_grid(self):
+        latlon = latlon_grid(bounds=((90, -90), (0, 360)), shape=(self.H, self.W))
+        self.latlon_torch = torch.tensor(np.stack(latlon, axis=0), dtype=torch.float32)
+
+    def _filename_to_index(self, filename):
+        """将 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 / self.time_step)
+
+    def __len__(self):
+        return self.total_samples
+
+    def __getitem__(self, idx):
+        year_idx = idx // self.samples_per_year
+        step_idx = idx % self.samples_per_year
+        year = self.used_years[year_idx]
+        files = self.files[year]
+
+        # ── invar: 从前一阶段模型输出的 npy 文件读取 ──────────
+        invar_list = []
+        for i in range(step_idx, step_idx + self.input_steps):
+            data = np.load(files[i])
+            data = np.squeeze(data)  # [C, H, W]
+            invar_list.append(data)
+
+        # ── outvar: 从年度 h5 文件读取真实 ERA5 标签 ──────────
+        outvar_list = []
+        time_index = []
+        h5_path = os.path.join(self.dataset_dir, 'data', f'{year}.h5')
+        with h5py.File(h5_path, "r") as f:
+            for i in range(step_idx + self.input_steps, step_idx + self.input_steps + self.output_steps):
+                fname = os.path.basename(files[i])[:-4]  # 去掉 .npy，得到 YYYYMMDDHH
+                t_idx = self._filename_to_index(fname)
+                data = f["fields"][t_idx]  # [C, H, W]
+                data = data[self.channel_indices]
+                outvar_list.append(data)
+                time_index.append(fname)
+
+        invar = np.stack(invar_list, axis=0)   # [T, C, H, W]
+        outvar = np.stack(outvar_list, axis=0) # [T, C, H, W]
+        invar = torch.as_tensor(invar, dtype=torch.float32)
+        outvar = torch.as_tensor(outvar, dtype=torch.float32)
+
+        if self.normalize:
+            invar  = (invar  - self.mu) / self.sd
+            outvar = (outvar - self.mu) / self.sd
+
+        # ── 太阳天顶角 ────────────────────────────────────────
+        start_time = datetime(year, 1, 1, tzinfo=pytz.utc)
+        timestamps = np.array([
+            (start_time + timedelta(hours=(step_idx + t) * self.time_step)).timestamp()
+            for t in range(self.output_steps)
+        ])
+        timestamps = torch.from_numpy(timestamps)
+        cos_zenith = cos_zenith_angle(timestamps, latlon=self.latlon_torch).float()
+
+        return invar.squeeze(0), outvar.squeeze(0), cos_zenith, step_idx, time_index

scripts/fake_data.py

@@ -0,0 +1,144 @@
+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 generate_fake_npy(result_dir, n_vars, years, dims):
+    """
+    为 short/medium 阶段生成假的模型输出 npy 文件（供 train_medium/train_long 输入）。
+    只创建一个真实的 npy 文件，其余使用 symlink。
+    """
+    T = dims["T"]
+    H, W = dims["H"], dims["W"]
+    time_step = dims["time_step"]
+
+    data_root = os.path.join(result_dir, "data")
+    os.makedirs(data_root, exist_ok=True)
+
+    real_year = years[0]
+    real_dir = os.path.join(data_root, str(real_year))
+    os.makedirs(real_dir, exist_ok=True)
+
+    timestamp = f'{real_year}010100'
+    real_path = os.path.join(real_dir, f"{timestamp}.npy")
+    real_data = np.zeros((n_vars, H, W), dtype=np.float32)
+    np.save(real_path, real_data)
+
+    # 当前年份其余时间步 symlink
+    for i in range(T):
+        ts = (np.datetime64(f"{real_year}-01-01T00") + np.timedelta64(i * time_step, "h")
+              ).astype(str).replace("-", "").replace("T", "")[:10]
+        path = os.path.join(real_dir, f"{ts}.npy")
+        if path != real_path and not os.path.exists(path):
+            os.symlink(f"{timestamp}.npy", path)
+
+    # 其他年份 symlink
+    for y in years[1:]:
+        y_dir = os.path.join(data_root, str(y))
+        os.makedirs(y_dir, exist_ok=True)
+        for i in range(T):
+            ts = (np.datetime64(f"{y}-01-01T00") + np.timedelta64(i * time_step, "h")
+                  ).astype(str).replace("-", "").replace("T", "")[:10]
+            path = os.path.join(y_dir, f"{ts}.npy")
+            if not os.path.exists(path):
+                rel = os.path.relpath(real_path, y_dir)
+                os.symlink(rel, path)
+
+    print(f"  ✅ Fake npy data generated → {result_dir}")
+
+
+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("/work/"):
+        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
+    n_vars   = len(atm_vars)
+
+    # 主 ERA5 数据（供 train_short 直接读取，归一化参数已内嵌进每年 h5）
+    generate_fake_h5(cfg_datapipe.dataset.data_dir, atm_vars, years, DATASET_DIMS)
+
+    # 前阶段推理输出：result/short 供 train_medium，result/medium 供 train_long
+    for stage in ['short', 'medium']:
+        generate_fake_npy(f'./result/{stage}', n_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,162 @@
+import torch
+import os
+import sys
+import glob
+import numpy as np
+import h5py
+from tqdm import tqdm
+from onescience.models.fuxi import Fuxi
+from onescience.utils.YParams import YParams
+
+
+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__":
+    if len(sys.argv) != 2:
+        print("Usage: input the mode: : short, medium, or long...")
+        sys.exit(1)
+
+    mode = sys.argv[1]
+    if mode not in ['short', 'medium', 'long']:
+        print(f'❌ ❌ Please input the mode: short, medium, or long...')
+        exit()
+
+    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")
+    cfg_data.dataloader.batch_size = 1
+    means, stds = get_stats(cfg_data.dataset.data_dir, cfg_data.dataset.channels)
+
+    if mode == 'short':
+        from onescience.datapipes.climate import ERA5Datapipe
+        train_datapipe = ERA5Datapipe(
+            dataset_dir=cfg_data.dataset.data_dir,
+            used_variables=cfg_data.dataset.channels,
+            used_years=cfg_data.dataset.train_time,
+            distributed=False,
+            input_steps=2,
+            batch_size=1,
+            num_workers=4,
+        )
+        train_dataloader, train_sampler = train_datapipe.get_dataloader("train")
+        val_datapipe = ERA5Datapipe(
+            dataset_dir=cfg_data.dataset.data_dir,
+            used_variables=cfg_data.dataset.channels,
+            used_years=cfg_data.dataset.val_time,
+            distributed=False,
+            input_steps=2,
+            batch_size=1,
+            num_workers=4,
+        )
+        val_dataloader, val_sampler = val_datapipe.get_dataloader("valid")
+        test_datapipe = ERA5Datapipe(
+            dataset_dir=cfg_data.dataset.data_dir,
+            used_variables=cfg_data.dataset.channels,
+            used_years=cfg_data.dataset.test_time,
+            distributed=False,
+            input_steps=2,
+            batch_size=1,
+            num_workers=4,
+        )
+        test_dataloader, _ = test_datapipe.get_dataloader("test")
+    else:
+        from data_loader import ERA5Datapipe
+        train_datapipe = ERA5Datapipe(
+            dataset_dir=cfg_data.dataset.data_dir,
+            used_variables=cfg_data.dataset.channels,
+            used_years=cfg_data.dataset.train_time,
+            pattern=mode,
+            distributed=False,
+            input_steps=2,
+            batch_size=1,
+            num_workers=4,
+        )
+        train_dataloader, train_sampler = train_datapipe.get_dataloader("train")
+        val_datapipe = ERA5Datapipe(
+            dataset_dir=cfg_data.dataset.data_dir,
+            used_variables=cfg_data.dataset.channels,
+            used_years=cfg_data.dataset.val_time,
+            pattern=mode,
+            distributed=False,
+            input_steps=2,
+            batch_size=1,
+            num_workers=4,
+        )
+        val_dataloader, val_sampler = val_datapipe.get_dataloader("valid")
+        test_datapipe = ERA5Datapipe(
+            dataset_dir=cfg_data.dataset.data_dir,
+            used_variables=cfg_data.dataset.channels,
+            used_years=cfg_data.dataset.test_time,
+            pattern=mode,
+            distributed=False,
+            input_steps=2,
+            batch_size=1,
+            num_workers=4,
+        )
+        test_dataloader, _ = test_datapipe.get_dataloader("test")
+
+    ckpt = torch.load(f"{cfg.checkpoint_dir}/model_{mode}_bak.pth", map_location="cuda:0")
+    model = Fuxi(img_size=cfg_data.dataset.img_size,
+                 patch_size=cfg.patch_size,
+                 in_chans=len(cfg_data.dataset.channels),
+                 out_chans=len(cfg_data.dataset.channels),
+                 embed_dim=cfg.embed_dim,
+                 num_groups=cfg.num_groups,
+                 num_heads=cfg.num_heads,
+                 window_size=cfg.window_size
+                 ).to("cuda:0")
+    model.load_state_dict(ckpt["model_state_dict"])
+
+    model.eval()
+    save_path = f'./result/{mode}/data/'
+    if mode != 'long':
+        with torch.no_grad():
+            print(f"📂 infer results will be generated to './result/{mode}/data/'")
+            for data in tqdm(train_dataloader, desc="Inferring trainset", unit="batch"):
+                invar = data[0].to("cuda:0", dtype=torch.float32)  # B, T, C, H, W
+                invar = invar.permute(0, 2, 1, 3, 4)  # B, C, T, H, W
+                filename = data[4][-1][0]
+                pred_var = model(invar).cpu().numpy()
+                pred_var = pred_var * stds + means
+                os.makedirs(f'{save_path}/{filename[:4]}', exist_ok=True)
+                np.save(f"{save_path}/{filename[:4]}/{filename}.npy", pred_var)
+
+        with torch.no_grad():
+            print(f"📂 infer results will be generated to './result/{mode}/data/'")
+            for data in tqdm(val_dataloader, desc="Inferring validset", unit="batch"):
+                invar = data[0].to("cuda:0", dtype=torch.float32)
+                invar = invar.permute(0, 2, 1, 3, 4)
+                filename = data[4][-1][0]
+                pred_var = model(invar).cpu().numpy()
+                pred_var = pred_var * stds + means
+                os.makedirs(f'{save_path}/{filename[:4]}', exist_ok=True)
+                np.save(f"{save_path}/{filename[:4]}/{filename}.npy", pred_var)
+
+    with torch.no_grad():
+        print(f"📂 infer results will be generated to './result/{mode}/data/'")
+        for data in tqdm(test_dataloader, desc="Inferring testset", unit="batch"):
+            invar = data[0].to("cuda:0", dtype=torch.float32)
+            invar = invar.permute(0, 2, 1, 3, 4)
+            filename = data[4][-1][0]
+            pred_var = model(invar).cpu().numpy()
+            pred_var = pred_var * stds + means
+            os.makedirs(f'{save_path}/{filename[:4]}', exist_ok=True)
+            np.save(f"{save_path}/{filename[:4]}/{filename}.npy", pred_var)

scripts/result.py

@@ -0,0 +1,214 @@
+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, test_years, mode):
+    """从新版 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 = []
+    for year in test_years:
+        result_dir = f'./result/{mode}/data/{year}/'
+        if os.path.exists(result_dir):
+            files = [f for f in os.listdir(result_dir) if f.endswith('.npy')]
+            total_files.extend(files)
+    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, mode):
+    channel_rmse = np.zeros(len(channel_indices))
+    channel_acc = np.zeros(len(channel_indices))
+    clim_mean = clim_mean[0, :, :, :]
+    if not os.path.exists(f'./result/{mode}_rmse.npy') or not os.path.exists(f'./result/{mode}_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/{mode}/data/{year}/{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(f'./result/{mode}_acc.npy', channel_acc)
+        np.save(f'./result/{mode}_rmse.npy', channel_rmse)
+
+
+def show_result(mode):
+    channel_rmse = np.load(f'./result/{mode}_rmse.npy')
+    channel_acc = np.load(f'./result/{mode}_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())
+
+    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__":
+    if len(sys.argv) != 2:
+        print("Usage: input the mode: : short, medium, or long...")
+        sys.exit(1)
+
+    mode = sys.argv[1]
+    if mode not in ['short', 'medium', 'long']:
+        print(f'❌ ❌ Please input the mode: short, medium, or long...')
+        exit()
+
+    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(f'./data/checkpoints/tr_{mode}_loss.npy')
+    valid_loss = np.load(f'./data/checkpoints/va_{mode}_loss.npy')
+    plot_loss(train_loss, valid_loss)
+
+    data_dir = cfg_data.dataset.data_dir
+    test_years = cfg_data.dataset.test_time
+    total_files, channel_indices, time_step = get_metadata(data_dir, cfg_data.dataset.channels, test_years, mode)
+
+    # 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, mode)
+    show_result(mode)
+
+    ##### 默认绘制 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/{mode}/data/{year}/{file}.npy').squeeze()
+        for i in range(len(selected_var)):
+            filename = f'./result/{mode}_{file}_{selected_var[i]}.png'
+            plot(label[channel_index[i]], pred[channel_index[i]], selected_var[i], filename)
+            print(f'✅plot {filename}')

scripts/train_long.py

@@ -0,0 +1,284 @@
+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
+
+from tqdm import tqdm
+from torch.nn.parallel import DistributedDataParallel
+from model.fuxi import Fuxi
+from scripts.data_loader import ERA5Datapipe
+from onescience.utils.YParams import YParams
+from onescience.metrics.climate.loss import LatitudeWeightedLoss
+from onescience.memory.checkpoint import replace_function
+
+from apex import optimizers
+
+
+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()
+    if not os.path.exists(f"{cfg.checkpoint_dir}/model_medium_bak.pth"):
+        if world_rank == 0:
+            print(f'❌❌The Fuxi medium model must be trained before this model.')
+        exit()
+    ## 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,
+        pattern='long',
+        distributed=dist.is_initialized(),
+        output_steps=2,
+        input_steps=2,
+        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,
+        pattern='long',
+        distributed=dist.is_initialized(),
+        output_steps=2,
+        input_steps=2,
+        batch_size=cfg_data.dataloader.batch_size,
+        num_workers=cfg_data.dataloader.num_workers
+    )
+    val_dataloader, val_sampler = datapipe.get_dataloader("valid")
+
+    ## Model init
+    model = Fuxi(img_size=cfg_data.dataset.img_size, 
+                 patch_size=cfg.patch_size, 
+                 in_chans=len(cfg_data.dataset.channels),
+                 out_chans=len(cfg_data.dataset.channels),
+                 embed_dim=cfg.embed_dim, 
+                 num_groups=cfg.num_groups, 
+                 num_heads=cfg.num_heads, 
+                 window_size=cfg.window_size
+                 ).to(local_rank)
+    optimizer = optimizers.FusedAdam(model.parameters(), lr=cfg.train_lr)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=5, mode="min")
+    loss_obj = LatitudeWeightedLoss(loss_type="l1", normalize=True).to(local_rank)
+
+    ## Train process init
+    os.makedirs(cfg.checkpoint_dir, exist_ok=True)
+    train_loss_file = f"{cfg.checkpoint_dir}/tr_long_loss.npy"
+    valid_loss_file = f"{cfg.checkpoint_dir}/va_long_loss.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 not os.path.exists(f"{cfg.checkpoint_dir}/model_medium_bak.pth"):
+        print('⚠️ ⚠️ Please train to get medium model first...')
+        exit()
+    
+    if os.path.exists(f"{cfg.checkpoint_dir}/model_long_bak.pth"):
+        if world_rank == 0:
+            print("\n\n")
+            print("-" * 50)
+            print(f"✅ There has a long-pattern model weight, load and continue training...")
+            print(f'If you want to finetune a new model, ensure there is no model_short_bak.pth file in {cfg.checkpoint_dir}')
+            print("-" * 50, "\n")
+        ckpt = torch.load(f"{cfg.checkpoint_dir}/model_long_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(f"{cfg.checkpoint_dir}/tr_long_loss.npy")
+        valid_losses = np.load(f"{cfg.checkpoint_dir}/va_long_loss.npy")
+    else:
+        if world_rank == 0:
+            print("\n\n")
+            print("-" * 50)
+            print(f"✅ Load medium model and continue to finetune...")
+            print("-" * 50, "\n")
+        ckpt = torch.load(f"{cfg.checkpoint_dir}/model_medium_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"])
+
+     ## 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(current_epoch, cfg.finetune_step):
+        if epoch > cfg.step_change_freq:
+            num_rollout_steps = epoch // cfg.step_change_freq + 2
+            if num_rollout_steps > 12: # Paper: 2~12 curriculum training schedule, then skip to 20.
+                num_rollout_steps = cfg.long_num_steps - cfg.medium_num_steps
+            if epoch % cfg.step_change_freq == 0 and world_rank == 0:
+                logger.info(f"⚠️ ⚠️ Switching to {num_rollout_steps}-step rollout!")
+            datapipe = ERA5Datapipe(
+                dataset_dir=cfg_data.dataset.data_dir,
+                used_variables=cfg_data.dataset.channels,
+                used_years=cfg_data.dataset.train_time,
+                pattern='long',
+                distributed=dist.is_initialized(),
+                output_steps=num_rollout_steps,
+                input_steps=2,
+                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,
+                pattern='long',
+                distributed=dist.is_initialized(),
+                output_steps=num_rollout_steps,
+                input_steps=2,
+                batch_size=cfg_data.dataloader.batch_size,
+                num_workers=cfg_data.dataloader.num_workers
+            )
+            val_dataloader, val_sampler = datapipe.get_dataloader("valid")
+        
+        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) # B, T, C, H, W
+            invar = invar.permute(0, 2, 1, 3, 4) # B, C, T, H, W
+            outvar = data[1].to(local_rank, dtype=torch.float32)
+            for t in range(outvar.shape[1]):
+                if t < outvar.shape[1] - 1:
+                    with torch.no_grad():
+                        outvar_pred = model(invar)
+                    # B, 70, 2, 721, 1440
+                    invar[:, :, 0] = invar[:, :, -1]
+                    invar[:, :, -1] = outvar_pred.detach()
+                else:
+                    with replace_function(model, ["cube_embedding", "u_transformer"], cfg.world_size > 1):
+                        outvar_pred = model(invar)
+                    loss = loss_obj(outvar_pred, outvar[:, t])
+            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].to(local_rank, dtype=torch.float32) # B, T, C, H, W
+                invar = invar.permute(0, 2, 1, 3, 4) # B, C, T, H, W
+                outvar = data[1].to(local_rank, dtype=torch.float32)
+                for t in range(outvar.shape[1]):
+                    outvar_pred = model(invar)
+                    # B, 70, 2, 721, 1440
+                    invar[:, :, 0] = invar[:, :, -1]
+                    invar[:, :, -1] = outvar_pred.detach()
+                loss = loss_obj(outvar_pred, outvar[:, -1])
+
+                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, epoch)
+            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, 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_long.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_long.pth {model_path}/model_long_bak.pth")
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    main()

scripts/train_medium.py

@@ -0,0 +1,285 @@
+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
+
+from tqdm import tqdm
+from torch.nn.parallel import DistributedDataParallel
+from model.fuxi import Fuxi
+from scripts.data_loader import ERA5Datapipe
+from onescience.utils.YParams import YParams
+from onescience.metrics.climate.loss import LatitudeWeightedLoss
+from onescience.memory.checkpoint import replace_function
+
+from apex import optimizers
+
+
+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()
+    if not os.path.exists(f"{cfg.checkpoint_dir}/model_short_bak.pth"):
+        if world_rank == 0:
+            print(f'❌❌The Fuxi short model must be trained before this model.')
+        exit()
+    ## 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,
+        pattern='medium',
+        distributed=dist.is_initialized(),
+        output_steps=2,
+        input_steps=2,
+        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,
+        pattern='medium',
+        distributed=dist.is_initialized(),
+        output_steps=2,
+        input_steps=2,
+        batch_size=cfg_data.dataloader.batch_size,
+        num_workers=cfg_data.dataloader.num_workers
+    )
+    val_dataloader, val_sampler = datapipe.get_dataloader("valid")
+
+    ## Model init
+    model = Fuxi(img_size=cfg_data.dataset.img_size, 
+                 patch_size=cfg.patch_size, 
+                 in_chans=len(cfg_data.dataset.channels),
+                 out_chans=len(cfg_data.dataset.channels),
+                 embed_dim=cfg.embed_dim, 
+                 num_groups=cfg.num_groups, 
+                 num_heads=cfg.num_heads, 
+                 window_size=cfg.window_size
+                 ).to(local_rank)
+    optimizer = optimizers.FusedAdam(model.parameters(), lr=cfg.train_lr)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=5, mode="min")
+    loss_obj = LatitudeWeightedLoss(loss_type="l1", normalize=True).to(local_rank)
+
+    ## Train process init
+    os.makedirs(cfg.checkpoint_dir, exist_ok=True)
+    train_loss_file = f"{cfg.checkpoint_dir}/tr_medium_loss.npy"
+    valid_loss_file = f"{cfg.checkpoint_dir}/va_medium_loss.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 
+
+    ## Load model weight if there exist well-trained model 
+    if not os.path.exists(f"{cfg.checkpoint_dir}/model_short_bak.pth"):
+        print('⚠️ ⚠️ Please train to get short model first...')
+        exit()
+    
+    if os.path.exists(f"{cfg.checkpoint_dir}/model_medium_bak.pth"):
+        if world_rank == 0:
+            print("\n\n")
+            print("-" * 50)
+            print(f"✅ There has a medium-pattern model weight, load and continue training...")
+            print(f'If you want to finetune a new model, ensure there is no model_short_bak.pth file in {cfg.checkpoint_dir}')
+            print("-" * 50, "\n")
+        ckpt = torch.load(f"{cfg.checkpoint_dir}/model_medium_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(f"{cfg.checkpoint_dir}/tr_medium_loss.npy")
+        valid_losses = np.load(f"{cfg.checkpoint_dir}/va_medium_loss.npy")
+    else:
+        if world_rank == 0:
+            print("\n\n")
+            print("-" * 50)
+            print(f"✅ Load short model and continue to finetune...")
+            print("-" * 50, "\n")
+        ckpt = torch.load(f"{cfg.checkpoint_dir}/model_short_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"])
+
+
+     ## 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(current_epoch, cfg.finetune_step):
+        if epoch % cfg.step_change_freq == 0:
+            num_rollout_steps = epoch // cfg.step_change_freq + 2
+            if num_rollout_steps > 12: # Paper: 2~12 curriculum training schedule, then skip to 20.
+                num_rollout_steps = cfg.medium_num_steps - cfg.short_num_steps
+            world_rank == 0 and logger.info(f"⚠️ ⚠️ Switching to {num_rollout_steps}-step rollout!")
+            datapipe = ERA5Datapipe(
+                dataset_dir=cfg_data.dataset.data_dir,
+                used_variables=cfg_data.dataset.channels,
+                used_years=cfg_data.dataset.train_time,
+                pattern='medium',
+                distributed=dist.is_initialized(),
+                output_steps=num_rollout_steps,
+                input_steps=2,
+                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,
+                pattern='medium',
+                distributed=dist.is_initialized(),
+                output_steps=num_rollout_steps,
+                input_steps=2,
+                batch_size=cfg_data.dataloader.batch_size,
+                num_workers=cfg_data.dataloader.num_workers
+            )
+            val_dataloader, val_sampler = datapipe.get_dataloader("valid")
+        
+        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) # B, T, C, H, W
+            invar = invar.permute(0, 2, 1, 3, 4) # B, C, T, H, W
+            outvar = data[1].to(local_rank, dtype=torch.float32)
+            for t in range(outvar.shape[1]):
+                if t < outvar.shape[1] - 1:
+                    with torch.no_grad():
+                        outvar_pred = model(invar)
+                    # B, 70, 2, 721, 1440
+                    invar[:, :, 0] = invar[:, :, -1]
+                    invar[:, :, -1] = outvar_pred.detach()
+                else:
+                    with replace_function(model, ["cube_embedding", "u_transformer"], cfg.world_size > 1):
+                        outvar_pred = model(invar)
+                    loss = loss_obj(outvar_pred, outvar[:, t])
+            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].to(local_rank, dtype=torch.float32) # B, T, C, H, W
+                invar = invar.permute(0, 2, 1, 3, 4) # B, C, T, H, W
+                outvar = data[1].to(local_rank, dtype=torch.float32)
+                for t in range(outvar.shape[1]):
+                    outvar_pred = model(invar)
+                    # B, 70, 2, 721, 1440
+                    invar[:, :, 0] = invar[:, :, -1]
+                    invar[:, :, -1] = outvar_pred.detach()
+                loss = loss_obj(outvar_pred, outvar[:, -1])
+
+                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, epoch)
+            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, 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_medium.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_medium.pth {model_path}/model_medium_bak.pth")
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    main()

scripts/train_short.py

@@ -0,0 +1,267 @@
+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
+
+from tqdm import tqdm
+from torch.nn.parallel import DistributedDataParallel
+from model.fuxi import Fuxi
+from onescience.datapipes.climate import ERA5Datapipe
+from onescience.utils.YParams import YParams
+from onescience.metrics.climate.loss import LatitudeWeightedLoss
+from onescience.memory.checkpoint import replace_function
+
+from apex import optimizers
+
+
+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(),
+        output_steps=2,
+        input_steps=2,
+        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(),
+        output_steps=2,
+        input_steps=2,
+        batch_size=cfg_data.dataloader.batch_size,
+        num_workers=cfg_data.dataloader.num_workers
+    )
+    val_dataloader, val_sampler = datapipe.get_dataloader("valid")
+
+    ## Model init
+    model = Fuxi(img_size=cfg_data.dataset.img_size, 
+                 patch_size=cfg.patch_size, 
+                 in_chans=len(cfg_data.dataset.channels),
+                 out_chans=len(cfg_data.dataset.channels),
+                 embed_dim=cfg.embed_dim, 
+                 num_groups=cfg.num_groups, 
+                 num_heads=cfg.num_heads, 
+                 window_size=cfg.window_size
+                 ).to(local_rank)
+    optimizer = optimizers.FusedAdam(model.parameters(), lr=cfg.train_lr)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=5, mode="min")
+    loss_obj = LatitudeWeightedLoss(loss_type="l1", normalize=True).to(local_rank)
+
+    ## Train process init
+    os.makedirs(cfg.checkpoint_dir, exist_ok=True)
+    train_loss_file = f"{cfg.checkpoint_dir}/tr_short_loss.npy"
+    valid_loss_file = f"{cfg.checkpoint_dir}/va_short_loss.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_short_bak.pth"):
+        if world_rank == 0:
+            print("\n\n")
+            print("-" * 50)
+            print(f"✅ There has a short-pattern model weight, load and continue training...")
+            print(f'If you want to finetune a new model, ensure there is no model_short_bak.pth file in {cfg.checkpoint_dir}')
+            print("-" * 50, "\n")
+        ckpt = torch.load(f"{cfg.checkpoint_dir}/model_short_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(f"{cfg.checkpoint_dir}/tr_short_loss.npy")
+        valid_losses = np.load(f"{cfg.checkpoint_dir}/va_short_loss.npy")
+    else:
+        if world_rank == 0:
+            print("\n\n")
+            print("-" * 50)
+            print(f"✅ No checkpoint found, training short model from scratch...")
+            print("-" * 50, "\n")
+
+
+     ## 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(current_epoch, cfg.finetune_step):
+        if epoch > cfg.step_change_freq:
+            num_rollout_steps = epoch // cfg.step_change_freq + 2
+            if num_rollout_steps > 12: # Paper: 2~12 curriculum training schedule, then skip to 20.
+                num_rollout_steps = cfg.short_num_steps
+            if epoch % cfg.step_change_freq == 0 and world_rank == 0:
+                logger.info(f"⚠️ ⚠️ Switching to {num_rollout_steps}-step rollout!")
+            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(),
+                output_steps=num_rollout_steps,
+                input_steps=2,
+              	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(),
+                output_steps=num_rollout_steps,
+                input_steps=2,
+              	batch_size=cfg_data.dataloader.batch_size,
+                num_workers=cfg_data.dataloader.num_workers
+            )
+            val_dataloader, val_sampler = datapipe.get_dataloader("valid")
+        
+        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) # B, T, C, H, W
+            invar = invar.permute(0, 2, 1, 3, 4) # B, C, T, H, W
+            outvar = data[1].to(local_rank, dtype=torch.float32)
+            for t in range(outvar.shape[1]):
+                if t < outvar.shape[1] - 1:
+                    with torch.no_grad():
+                        outvar_pred = model(invar)
+                    # B, 70, 2, 721, 1440
+                    invar[:, :, 0] = invar[:, :, -1]
+                    invar[:, :, -1] = outvar_pred.detach()
+                else:
+                    with replace_function(model, ["cube_embedding", "u_transformer"], cfg.world_size > 1):
+                        outvar_pred = model(invar)
+                    loss = loss_obj(outvar_pred, outvar[:, t])
+
+            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].to(local_rank, dtype=torch.float32) # B, T, C, H, W
+                invar = invar.permute(0, 2, 1, 3, 4) # B, C, T, H, W
+                outvar = data[1].to(local_rank, dtype=torch.float32)
+                for t in range(outvar.shape[1]):
+                    outvar_pred = model(invar)
+                    # B, 70, 2, 721, 1440
+                    invar[:, :, 0] = invar[:, :, -1]
+                    invar[:, :, -1] = outvar_pred.detach()
+                loss = loss_obj(outvar_pred, outvar[:, -1])
+
+                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, epoch)
+            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, 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_short.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_short.pth {model_path}/model_short_bak.pth")
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    main()