Upload folder using huggingface_hub

README.md

@@ -0,0 +1,185 @@
+---
+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;">GraphCast</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+GraphCast is a global medium-range weather forecast model developed by the Google DeepMind team, with its core paper published in the top-tier international journal *Science*.
+
+Paper: GraphCast: Learning skillful medium-range global weather forecasting
+
+https://arxiv.org/abs/2212.12794
+
+# Model Description
+
+GraphCast is a global medium-range weather forecast model built on a Graph Neural Network (GNN). It is trained on the ERA5 global atmospheric reanalysis dataset (1979–2017) provided by ECMWF.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Global Weather Forecast Research | Train a GraphCast-style GNN forecast model using annual ERA5 HDF5 data. |
+| Local Quick Validation | Use synthetic data to verify data loading, auxiliary file generation, training entry points, and result scripts. |
+| 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/GraphCast --local_dir ./GraphCast
+cd GraphCast
+```
+
+### 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
+```
+
+### Generate Auxiliary Files
+
+```bash
+python scripts/get_data_json.py
+python scripts/compute_time_diff_std.py
+```
+
+Generated files:
+
+- `data.json`
+- `time_diff_std.npy`
+
+### 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 outputs:
+
+```text
+data/checkpoints/model_bak.pth
+data/checkpoints/trloss.npy
+```
+
+### Training Weights
+
+This repository provides weights trained on ERA5 data from 1979 to 2017 in the `weights/` folder. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Fine-tuning
+
+Before fine-tuning, you must first complete training and generate `data/checkpoints/model_bak.pth`.
+
+```bash
+python scripts/finetune.py
+```
+
+Fine-tuning outputs:
+
+```text
+data/checkpoints/model_finetune_bak.pth
+data/checkpoints/ft_trloss.npy
+```
+
+### Inference
+
+Inference reads `data/checkpoints/model_finetune_bak.pth` by default:
+
+```bash
+python scripts/inference.py
+```
+
+Prediction results are output to:
+
+```text
+result/output/
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Output contents include:
+
+- `result/rmse.npy`
+- `result/acc.npy`
+- `result/loss.png`
+- Forecast comparison plots for specified dates and variables
+
+# 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
+
+- Apache License 2.0. The code is open source, permitting both commercial and non-commercial use.
+- The weights are permitted for non-commercial use only.

conf/config.yaml

@@ -0,0 +1,186 @@
+model:
+  processor_layers: 16          
+  hidden_dim: 512               
+  mesh_level: 6                 
+  multimesh: true               
+  num_channels_climate: 227
+  use_cos_zenith: true         
+  use_time_of_year_index: true 
+  num_history: 0               
+  num_channels_static: 5       
+  num_steps: 1
+  checkpoint_encoder: true      
+  checkpoint_decoder: false     
+  full_bf16: true              
+  processor_type: MessagePassing
+  khop_neighbors: 32            
+  num_attention_heads: 4        
+  norm_type: LayerNorm          
+  segments: 1                   
+  force_single_checkpoint: false
+  checkpoint_processor: false   
+  force_single_checkpoint_finetune: false
+  checkpoint_encoder_finetune: true      
+  checkpoint_processor_finetune: true    
+  checkpoint_decoder_finetune: true      
+  concat_trick: true         
+  cugraphops_encoder: false  
+  cugraphops_processor: false
+  cugraphops_decoder: false  
+  recompute_activation: false
+  use_apex: true             
+  patience: 50
+  start_epoch: 0
+  max_epoch: 9999
+  num_val_steps: 8      
+  dt: 6.0                
+  grad_clip_norm: 32.0  
+  lr: 1e-3               
+  lr_step3: 3e-7        
+  num_iters_step1: 1000  
+  num_iters_step2: 299000
+  num_iters_step3: 11000 
+  step_change_freq: 1000
+  val_freq: 5           
+  checkpoint_dir : "./data/checkpoints"
+  
+# 整个数据读取流程
+datapipe:
+  name: "ERA5"
+  task: "weather_forecasting"
+
+  # dataset设定
+  dataset:
+    type: "hdf5"
+    data_dir: './data/' # "$ONESCIENCE_DATASETS_DIR/ERA5/newh5/" 
+    dataset_metadata_path: './data.json'  # Path to the dataset metadata, containing channel names.
+    time_diff_std_path: './time_diff_std.npy'  # Path to the .npy file with standard deviation of normalized per-variable per-pressure level time differences.
+    train_time: [1951, 1952]
+    val_time: [1953]
+    test_time: [1954]
+    img_size: [721, 1440]
+    verbose: true
+    cache: false
+    # vars
+    channels: ['10m_u_component_of_wind', '10m_v_component_of_wind', '2m_temperature', 'mean_sea_level_pressure', 'total_precipitation',
+
+               'geopotential_1', 'geopotential_2', 'geopotential_3', 'geopotential_5', 'geopotential_7', 
+               'geopotential_10', 'geopotential_20', 'geopotential_30', 'geopotential_50', 'geopotential_70', 
+               'geopotential_100', 'geopotential_125', 'geopotential_150', 'geopotential_175', 'geopotential_200', 
+               'geopotential_225', 'geopotential_250', 'geopotential_300', 'geopotential_350', 'geopotential_400', 
+               'geopotential_450', 'geopotential_500', 'geopotential_550', 'geopotential_600', 'geopotential_650', 
+               'geopotential_700', 'geopotential_750', 'geopotential_775', 'geopotential_800', 'geopotential_825', 
+               'geopotential_850', 'geopotential_875', 'geopotential_900', 'geopotential_925', 'geopotential_950', 
+               'geopotential_975', 'geopotential_1000',
+
+               'relative_humidity_1', 'relative_humidity_2', 'relative_humidity_3', 'relative_humidity_5', 'relative_humidity_7', 
+               'relative_humidity_10', 'relative_humidity_20', 'relative_humidity_30', 'relative_humidity_50', 'relative_humidity_70', 
+               'relative_humidity_100', 'relative_humidity_125', 'relative_humidity_150', 'relative_humidity_175', 'relative_humidity_200', 
+               'relative_humidity_225', 'relative_humidity_250', 'relative_humidity_300', 'relative_humidity_350', 'relative_humidity_400', 
+               'relative_humidity_450', 'relative_humidity_500', 'relative_humidity_550', 'relative_humidity_600', 'relative_humidity_650', 
+               'relative_humidity_700', 'relative_humidity_750', 'relative_humidity_775', 'relative_humidity_800', 'relative_humidity_825', 
+               'relative_humidity_850', 'relative_humidity_875', 'relative_humidity_900', 'relative_humidity_925', 'relative_humidity_950', 
+               'relative_humidity_975', 'relative_humidity_1000',
+
+               'u_component_of_wind_1', 'u_component_of_wind_2', 'u_component_of_wind_3', 'u_component_of_wind_5', 'u_component_of_wind_7', 
+               'u_component_of_wind_10', 'u_component_of_wind_20', 'u_component_of_wind_30', 'u_component_of_wind_50', 'u_component_of_wind_70', 
+               'u_component_of_wind_100', 'u_component_of_wind_125', 'u_component_of_wind_150', 'u_component_of_wind_175', 'u_component_of_wind_200', 
+               'u_component_of_wind_225', 'u_component_of_wind_250', 'u_component_of_wind_300', 'u_component_of_wind_350', 'u_component_of_wind_400', 
+               'u_component_of_wind_450', 'u_component_of_wind_500', 'u_component_of_wind_550', 'u_component_of_wind_600', 'u_component_of_wind_650', 
+               'u_component_of_wind_700', 'u_component_of_wind_750', 'u_component_of_wind_775', 'u_component_of_wind_800', 'u_component_of_wind_825', 
+               'u_component_of_wind_850', 'u_component_of_wind_875', 'u_component_of_wind_900', 'u_component_of_wind_925', 'u_component_of_wind_950', 
+               'u_component_of_wind_975', 'u_component_of_wind_1000',
+
+               'v_component_of_wind_1', 'v_component_of_wind_2', 'v_component_of_wind_3', 'v_component_of_wind_5', 'v_component_of_wind_7', 
+               'v_component_of_wind_10', 'v_component_of_wind_20', 'v_component_of_wind_30', 'v_component_of_wind_50', 'v_component_of_wind_70', 
+               'v_component_of_wind_100', 'v_component_of_wind_125', 'v_component_of_wind_150', 'v_component_of_wind_175', 'v_component_of_wind_200', 
+               'v_component_of_wind_225', 'v_component_of_wind_250', 'v_component_of_wind_300', 'v_component_of_wind_350', 'v_component_of_wind_400', 
+               'v_component_of_wind_450', 'v_component_of_wind_500', 'v_component_of_wind_550', 'v_component_of_wind_600', 'v_component_of_wind_650', 
+               'v_component_of_wind_700', 'v_component_of_wind_750', 'v_component_of_wind_775', 'v_component_of_wind_800', 'v_component_of_wind_825', 
+               'v_component_of_wind_850', 'v_component_of_wind_875', 'v_component_of_wind_900', 'v_component_of_wind_925', 'v_component_of_wind_950', 
+               'v_component_of_wind_975', 'v_component_of_wind_1000',
+
+               'temperature_1', 'temperature_2', 'temperature_3', 'temperature_5', 'temperature_7', 
+               'temperature_10', 'temperature_20', 'temperature_30', 'temperature_50', 'temperature_70', 
+               'temperature_100', 'temperature_125', 'temperature_150', 'temperature_175', 'temperature_200', 
+               'temperature_225', 'temperature_250', 'temperature_300', 'temperature_350', 'temperature_400', 
+               'temperature_450', 'temperature_500', 'temperature_550', 'temperature_600', 'temperature_650', 
+               'temperature_700', 'temperature_750', 'temperature_775', 'temperature_800', 'temperature_825', 
+               'temperature_850', 'temperature_875', 'temperature_900', 'temperature_925', 'temperature_950', 
+               'temperature_975', 'temperature_1000',
+
+               'vertical_wind_speed_1', 'vertical_wind_speed_2', 'vertical_wind_speed_3', 'vertical_wind_speed_5', 'vertical_wind_speed_7', 
+               'vertical_wind_speed_10', 'vertical_wind_speed_20', 'vertical_wind_speed_30', 'vertical_wind_speed_50', 'vertical_wind_speed_70', 
+               'vertical_wind_speed_100', 'vertical_wind_speed_125', 'vertical_wind_speed_150', 'vertical_wind_speed_175', 'vertical_wind_speed_200', 
+               'vertical_wind_speed_225', 'vertical_wind_speed_250', 'vertical_wind_speed_300', 'vertical_wind_speed_350', 'vertical_wind_speed_400', 
+               'vertical_wind_speed_450', 'vertical_wind_speed_500', 'vertical_wind_speed_550', 'vertical_wind_speed_600', 'vertical_wind_speed_650', 
+               'vertical_wind_speed_700', 'vertical_wind_speed_750', 'vertical_wind_speed_775', 'vertical_wind_speed_800', 'vertical_wind_speed_825', 
+               'vertical_wind_speed_850', 'vertical_wind_speed_875', 'vertical_wind_speed_900', 'vertical_wind_speed_925', 'vertical_wind_speed_950', 
+               'vertical_wind_speed_975', 'vertical_wind_speed_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/graph_cast_net.py

@@ -0,0 +1,805 @@
+import logging
+import warnings
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+from torch import Tensor
+from torch import nn
+try:
+    from typing import Self
+except ImportError:
+    # for Python versions < 3.11
+    from typing_extensions import Self
+
+from onescience.modules.decoder.mesh_graph_decoder import MeshGraphDecoder
+from onescience.modules.embedding.graphcast_embedder import (
+    GraphCastDecoderEmbedder,
+    GraphCastEncoderEmbedder,
+)
+from onescience.modules.encoder.mesh_graph_encoder import MeshGraphEncoder
+from onescience.modules.mlp.mesh_graph_mlp import MeshGraphMLP
+from onescience.modules.utils.gnnlayer_utils import CuGraphCSC, set_checkpoint_fn
+from onescience.modules.layer.activations import get_activation
+from onescience.modules.utils.graphcast.graph import Graph
+
+from onescience.models.meta import ModelMetaData
+
+from .graph_cast_processor import (
+    GraphCastProcessor,
+    GraphCastProcessorGraphTransformer,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def get_lat_lon_partition_separators(partition_size: int):
+    """Utility Function to get separation intervals for lat-lon
+    grid for partition_sizes of interest.
+
+    Parameters
+    ----------
+    partition_size : int
+        size of graph partition
+    """
+
+    def _divide(num_lat_chunks: int, num_lon_chunks: int):
+        # divide lat-lon grid into equally-sizes chunks along both latitude and longitude
+        if (num_lon_chunks * num_lat_chunks) != partition_size:
+            raise ValueError(
+                "Can't divide lat-lon grid into grid {num_lat_chunks} x {num_lon_chunks} chunks for partition_size={partition_size}."
+            )
+        # divide latitutude into num_lat_chunks of size 180 / num_lat_chunks
+        # divide longitude into chunks of size 360 / (partition_size / num_lat_chunks)
+        lat_bin_width = 180.0 / num_lat_chunks
+        lon_bin_width = 360.0 / num_lon_chunks
+
+        lat_ranges = []
+        lon_ranges = []
+
+        for p_lat in range(num_lat_chunks):
+            for p_lon in range(num_lon_chunks):
+                lat_ranges += [
+                    (lat_bin_width * p_lat - 90.0, lat_bin_width * (p_lat + 1) - 90.0)
+                ]
+                lon_ranges += [
+                    (lon_bin_width * p_lon - 180.0, lon_bin_width * (p_lon + 1) - 180.0)
+                ]
+
+        lat_ranges[-1] = (lat_ranges[-1][0], None)
+        lon_ranges[-1] = (lon_ranges[-1][0], None)
+
+        return lat_ranges, lon_ranges
+
+    # use two closest factors of partition_size
+    lat_chunks, lon_chunks, i = 1, partition_size, 0
+    while lat_chunks < lon_chunks:
+        i += 1
+        if partition_size % i == 0:
+            lat_chunks = i
+            lon_chunks = partition_size // lat_chunks
+
+    lat_ranges, lon_ranges = _divide(lat_chunks, lon_chunks)
+
+    # mainly for debugging
+    if (lat_ranges is None) or (lon_ranges is None):
+        raise ValueError("unexpected error, abort")
+
+    min_seps = []
+    max_seps = []
+
+    for i in range(partition_size):
+        lat = lat_ranges[i]
+        lon = lon_ranges[i]
+        min_seps.append([lat[0], lon[0]])
+        max_seps.append([lat[1], lon[1]])
+
+    return min_seps, max_seps
+
+
+@dataclass
+class MetaData(ModelMetaData):
+    name: str = "GraphCastNet"
+    # Optimization
+    jit: bool = False
+    cuda_graphs: bool = False
+    amp_cpu: bool = False
+    amp_gpu: bool = True
+    torch_fx: bool = False
+    # Data type
+    bf16: bool = True
+    # Inference
+    onnx: bool = False
+    # Physics informed
+    func_torch: bool = False
+    auto_grad: bool = False
+
+
+class GraphCastNet(nn.Module):
+    def __init__(
+        self,
+        mesh_level: Optional[int] = 6,
+        multimesh_level: Optional[int] = None,
+        multimesh: bool = True,
+        input_res: tuple = (721, 1440),
+        input_dim_grid_nodes: int = 237,
+        input_dim_mesh_nodes: int = 3,
+        input_dim_edges: int = 4,
+        output_dim_grid_nodes: int = 227,
+        processor_type: str = "MessagePassing",
+        khop_neighbors: int = 32,
+        num_attention_heads: int = 4,
+        processor_layers: int = 16,
+        hidden_layers: int = 1,
+        hidden_dim: int = 512,
+        aggregation: str = "sum",
+        activation_fn: str = "silu",
+        norm_type: str = "LayerNorm",
+        use_cugraphops_encoder: bool = False,
+        use_cugraphops_processor: bool = False,
+        use_cugraphops_decoder: bool = False,
+        do_concat_trick: bool = False,
+        recompute_activation: bool = False,
+        partition_size: int = 1,
+        partition_group_name: Optional[str] = None,
+        use_lat_lon_partitioning: bool = False,
+        expect_partitioned_input: bool = False,
+        global_features_on_rank_0: bool = False,
+        produce_aggregated_output: bool = True,
+        produce_aggregated_output_on_all_ranks: bool = True,
+    ):
+        super().__init__()
+
+        # 'multimesh_level' deprecation handling
+        if multimesh_level is not None:
+            warnings.warn(
+                "'multimesh_level' is deprecated and will be removed in a future version. Use 'mesh_level' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            mesh_level = multimesh_level
+
+        self.processor_type = processor_type
+        if self.processor_type == "MessagePassing":
+            khop_neighbors = 0
+        self.is_distributed = False
+        if partition_size > 1:
+            self.is_distributed = True
+        self.expect_partitioned_input = expect_partitioned_input
+        self.global_features_on_rank_0 = global_features_on_rank_0
+        self.produce_aggregated_output = produce_aggregated_output
+        self.produce_aggregated_output_on_all_ranks = (
+            produce_aggregated_output_on_all_ranks
+        )
+        self.partition_group_name = partition_group_name
+
+        # create the lat_lon_grid
+        self.latitudes = torch.linspace(-90, 90, steps=input_res[0])
+        self.longitudes = torch.linspace(-180, 180, steps=input_res[1] + 1)[1:]
+        self.lat_lon_grid = torch.stack(
+            torch.meshgrid(self.latitudes, self.longitudes, indexing="ij"), dim=-1
+        )
+
+        # Set activation function
+        activation_fn = get_activation(activation_fn)
+
+        # construct the graph
+        self.graph = Graph(self.lat_lon_grid, mesh_level, multimesh, khop_neighbors)
+
+        self.mesh_graph, self.attn_mask = self.graph.create_mesh_graph(verbose=False)
+        self.g2m_graph = self.graph.create_g2m_graph(verbose=False)
+        self.m2g_graph = self.graph.create_m2g_graph(verbose=False)
+
+        self.g2m_edata = self.g2m_graph.edata["x"]
+        self.m2g_edata = self.m2g_graph.edata["x"]
+        self.mesh_ndata = self.mesh_graph.ndata["x"]
+        if self.processor_type == "MessagePassing":
+            self.mesh_edata = self.mesh_graph.edata["x"]
+        elif self.processor_type == "GraphTransformer":
+            # Dummy tensor to avoid breaking the API
+            self.mesh_edata = torch.zeros((1, input_dim_edges))
+        else:
+            raise ValueError(f"Invalid processor type {processor_type}")
+
+        if use_cugraphops_encoder or self.is_distributed:
+            kwargs = {}
+            if use_lat_lon_partitioning:
+                min_seps, max_seps = get_lat_lon_partition_separators(partition_size)
+                kwargs = {
+                    "src_coordinates": self.g2m_graph.srcdata["lat_lon"],
+                    "dst_coordinates": self.g2m_graph.dstdata["lat_lon"],
+                    "coordinate_separators_min": min_seps,
+                    "coordinate_separators_max": max_seps,
+                }
+            self.g2m_graph, edge_perm = CuGraphCSC.from_dgl(
+                graph=self.g2m_graph,
+                partition_size=partition_size,
+                partition_group_name=partition_group_name,
+                partition_by_bbox=use_lat_lon_partitioning,
+                **kwargs,
+            )
+            self.g2m_edata = self.g2m_edata[edge_perm]
+
+            if self.is_distributed:
+                self.g2m_edata = self.g2m_graph.get_edge_features_in_partition(
+                    self.g2m_edata
+                )
+
+        if use_cugraphops_decoder or self.is_distributed:
+            kwargs = {}
+            if use_lat_lon_partitioning:
+                min_seps, max_seps = get_lat_lon_partition_separators(partition_size)
+                kwargs = {
+                    "src_coordinates": self.m2g_graph.srcdata["lat_lon"],
+                    "dst_coordinates": self.m2g_graph.dstdata["lat_lon"],
+                    "coordinate_separators_min": min_seps,
+                    "coordinate_separators_max": max_seps,
+                }
+
+            self.m2g_graph, edge_perm = CuGraphCSC.from_dgl(
+                graph=self.m2g_graph,
+                partition_size=partition_size,
+                partition_group_name=partition_group_name,
+                partition_by_bbox=use_lat_lon_partitioning,
+                **kwargs,
+            )
+            self.m2g_edata = self.m2g_edata[edge_perm]
+
+            if self.is_distributed:
+                self.m2g_edata = self.m2g_graph.get_edge_features_in_partition(
+                    self.m2g_edata
+                )
+
+        if use_cugraphops_processor or self.is_distributed:
+            kwargs = {}
+            if use_lat_lon_partitioning:
+                min_seps, max_seps = get_lat_lon_partition_separators(partition_size)
+                kwargs = {
+                    "src_coordinates": self.mesh_graph.ndata["lat_lon"],
+                    "dst_coordinates": self.mesh_graph.ndata["lat_lon"],
+                    "coordinate_separators_min": min_seps,
+                    "coordinate_separators_max": max_seps,
+                }
+
+            self.mesh_graph, edge_perm = CuGraphCSC.from_dgl(
+                graph=self.mesh_graph,
+                partition_size=partition_size,
+                partition_group_name=partition_group_name,
+                partition_by_bbox=use_lat_lon_partitioning,
+                **kwargs,
+            )
+            self.mesh_edata = self.mesh_edata[edge_perm]
+            if self.is_distributed:
+                self.mesh_edata = self.mesh_graph.get_edge_features_in_partition(
+                    self.mesh_edata
+                )
+                self.mesh_ndata = self.mesh_graph.get_dst_node_features_in_partition(
+                    self.mesh_ndata
+                )
+
+        self.input_dim_grid_nodes = input_dim_grid_nodes
+        self.output_dim_grid_nodes = output_dim_grid_nodes
+        self.input_res = input_res
+
+        # by default: don't checkpoint at all
+        self.model_checkpoint_fn = set_checkpoint_fn(False)
+        self.encoder_checkpoint_fn = set_checkpoint_fn(False)
+        self.decoder_checkpoint_fn = set_checkpoint_fn(False)
+
+        # initial feature embedder
+        self.encoder_embedder = GraphCastEncoderEmbedder(
+            input_dim_grid_nodes=input_dim_grid_nodes,
+            input_dim_mesh_nodes=input_dim_mesh_nodes,
+            input_dim_edges=input_dim_edges,
+            output_dim=hidden_dim,
+            hidden_dim=hidden_dim,
+            hidden_layers=hidden_layers,
+            activation_fn=activation_fn,
+            norm_type=norm_type,
+            recompute_activation=recompute_activation,
+        )
+        self.decoder_embedder = GraphCastDecoderEmbedder(
+            input_dim_edges=input_dim_edges,
+            output_dim=hidden_dim,
+            hidden_dim=hidden_dim,
+            hidden_layers=hidden_layers,
+            activation_fn=activation_fn,
+            norm_type=norm_type,
+            recompute_activation=recompute_activation,
+        )
+
+        # grid2mesh encoder
+        self.encoder = MeshGraphEncoder(
+            aggregation=aggregation,
+            input_dim_src_nodes=hidden_dim,
+            input_dim_dst_nodes=hidden_dim,
+            input_dim_edges=hidden_dim,
+            output_dim_src_nodes=hidden_dim,
+            output_dim_dst_nodes=hidden_dim,
+            output_dim_edges=hidden_dim,
+            hidden_dim=hidden_dim,
+            hidden_layers=hidden_layers,
+            activation_fn=activation_fn,
+            norm_type=norm_type,
+            do_concat_trick=do_concat_trick,
+            recompute_activation=recompute_activation,
+        )
+
+        # icosahedron processor
+        if processor_layers <= 2:
+            raise ValueError("Expected at least 3 processor layers")
+        if processor_type == "MessagePassing":
+            self.processor_encoder = GraphCastProcessor(
+                aggregation=aggregation,
+                processor_layers=1,
+                input_dim_nodes=hidden_dim,
+                input_dim_edges=hidden_dim,
+                hidden_dim=hidden_dim,
+                hidden_layers=hidden_layers,
+                activation_fn=activation_fn,
+                norm_type=norm_type,
+                do_concat_trick=do_concat_trick,
+                recompute_activation=recompute_activation,
+            )
+            self.processor = GraphCastProcessor(
+                aggregation=aggregation,
+                processor_layers=processor_layers - 2,
+                input_dim_nodes=hidden_dim,
+                input_dim_edges=hidden_dim,
+                hidden_dim=hidden_dim,
+                hidden_layers=hidden_layers,
+                activation_fn=activation_fn,
+                norm_type=norm_type,
+                do_concat_trick=do_concat_trick,
+                recompute_activation=recompute_activation,
+            )
+            self.processor_decoder = GraphCastProcessor(
+                aggregation=aggregation,
+                processor_layers=1,
+                input_dim_nodes=hidden_dim,
+                input_dim_edges=hidden_dim,
+                hidden_dim=hidden_dim,
+                hidden_layers=hidden_layers,
+                activation_fn=activation_fn,
+                norm_type=norm_type,
+                do_concat_trick=do_concat_trick,
+                recompute_activation=recompute_activation,
+            )
+        else:
+            self.processor_encoder = torch.nn.Identity()
+            self.processor = GraphCastProcessorGraphTransformer(
+                attention_mask=self.attn_mask,
+                num_attention_heads=num_attention_heads,
+                processor_layers=processor_layers,
+                input_dim_nodes=hidden_dim,
+                hidden_dim=hidden_dim,
+            )
+            self.processor_decoder = torch.nn.Identity()
+
+        # mesh2grid decoder
+        self.decoder = MeshGraphDecoder(
+            aggregation=aggregation,
+            input_dim_src_nodes=hidden_dim,
+            input_dim_dst_nodes=hidden_dim,
+            input_dim_edges=hidden_dim,
+            output_dim_dst_nodes=hidden_dim,
+            output_dim_edges=hidden_dim,
+            hidden_dim=hidden_dim,
+            hidden_layers=hidden_layers,
+            activation_fn=activation_fn,
+            norm_type=norm_type,
+            do_concat_trick=do_concat_trick,
+            recompute_activation=recompute_activation,
+        )
+
+        # final MLP
+        self.finale = MeshGraphMLP(
+            input_dim=hidden_dim,
+            output_dim=output_dim_grid_nodes,
+            hidden_dim=hidden_dim,
+            hidden_layers=hidden_layers,
+            activation_fn=activation_fn,
+            norm_type=None,
+            recompute_activation=recompute_activation,
+        )
+
+    def set_checkpoint_model(self, checkpoint_flag: bool):
+        """Sets checkpoint function for the entire model.
+
+        This function returns the appropriate checkpoint function based on the
+        provided `checkpoint_flag` flag. If `checkpoint_flag` is True, the
+        function returns the checkpoint function from PyTorch's
+        `torch.utils.checkpoint`. In this case, all the other gradient checkpoitings
+        will be disabled. Otherwise, it returns an identity function
+        that simply passes the inputs through the given layer.
+
+        Parameters
+        ----------
+        checkpoint_flag : bool
+            Whether to use checkpointing for gradient computation. Checkpointing
+            can reduce memory usage during backpropagation at the cost of
+            increased computation time.
+
+        Returns
+        -------
+        Callable
+            The selected checkpoint function to use for gradient computation.
+        """
+        # force a single checkpoint for the whole model
+        self.model_checkpoint_fn = set_checkpoint_fn(checkpoint_flag)
+        if checkpoint_flag:
+            self.processor.set_checkpoint_segments(-1)
+            self.encoder_checkpoint_fn = set_checkpoint_fn(False)
+            self.decoder_checkpoint_fn = set_checkpoint_fn(False)
+
+    def set_checkpoint_processor(self, checkpoint_segments: int):
+        """Sets checkpoint function for the processor excluding the first and last
+        layers.
+
+        This function returns the appropriate checkpoint function based on the
+        provided `checkpoint_segments` flag. If `checkpoint_segments` is positive,
+        the function returns the checkpoint function from PyTorch's
+        `torch.utils.checkpoint`, with number of checkpointing segments equal to
+        `checkpoint_segments`. Otherwise, it returns an identity function
+        that simply passes the inputs through the given layer.
+
+        Parameters
+        ----------
+        checkpoint_segments : int
+            Number of checkpointing segments for gradient computation. Checkpointing
+            can reduce memory usage during backpropagation at the cost of
+            increased computation time.
+
+        Returns
+        -------
+        Callable
+            The selected checkpoint function to use for gradient computation.
+        """
+        self.processor.set_checkpoint_segments(checkpoint_segments)
+
+    def set_checkpoint_encoder(self, checkpoint_flag: bool):
+        """Sets checkpoint function for the embedder, encoder, and the first of
+        the processor.
+
+        This function returns the appropriate checkpoint function based on the
+        provided `checkpoint_flag` flag. If `checkpoint_flag` is True, the
+        function returns the checkpoint function from PyTorch's
+        `torch.utils.checkpoint`. Otherwise, it returns an identity function
+        that simply passes the inputs through the given layer.
+
+        Parameters
+        ----------
+        checkpoint_flag : bool
+            Whether to use checkpointing for gradient computation. Checkpointing
+            can reduce memory usage during backpropagation at the cost of
+            increased computation time.
+
+        Returns
+        -------
+        Callable
+            The selected checkpoint function to use for gradient computation.
+        """
+        self.encoder_checkpoint_fn = set_checkpoint_fn(checkpoint_flag)
+
+    def set_checkpoint_decoder(self, checkpoint_flag: bool):
+        """Sets checkpoint function for the last layer of the processor, the decoder,
+        and the final MLP.
+
+        This function returns the appropriate checkpoint function based on the
+        provided `checkpoint_flag` flag. If `checkpoint_flag` is True, the
+        function returns the checkpoint function from PyTorch's
+        `torch.utils.checkpoint`. Otherwise, it returns an identity function
+        that simply passes the inputs through the given layer.
+
+        Parameters
+        ----------
+        checkpoint_flag : bool
+            Whether to use checkpointing for gradient computation. Checkpointing
+            can reduce memory usage during backpropagation at the cost of
+            increased computation time.
+
+        Returns
+        -------
+        Callable
+            The selected checkpoint function to use for gradient computation.
+        """
+        self.decoder_checkpoint_fn = set_checkpoint_fn(checkpoint_flag)
+
+    def encoder_forward(
+        self,
+        grid_nfeat: Tensor,
+    ) -> Tensor:
+        """Forward method for the embedder, encoder, and the first of the processor.
+
+        Parameters
+        ----------
+        grid_nfeat : Tensor
+            Node features for the latitude-longitude grid.
+
+        Returns
+        -------
+        mesh_efeat_processed: Tensor
+            Processed edge features for the multimesh.
+        mesh_nfeat_processed: Tensor
+            Processed node features for the multimesh.
+        grid_nfeat_encoded: Tensor
+            Encoded node features for the latitude-longitude grid.
+        """
+
+        # embedd graph features
+        (
+            grid_nfeat_embedded,
+            mesh_nfeat_embedded,
+            g2m_efeat_embedded,
+            mesh_efeat_embedded,
+        ) = self.encoder_embedder(
+            grid_nfeat,
+            self.mesh_ndata,
+            self.g2m_edata,
+            self.mesh_edata
+        )
+
+        # encode lat/lon to multimesh
+        grid_nfeat_encoded, mesh_nfeat_encoded = self.encoder(
+            g2m_efeat_embedded,
+            grid_nfeat_embedded,
+            mesh_nfeat_embedded,
+            self.g2m_graph,
+        )
+
+        # process multimesh graph
+        if self.processor_type == "MessagePassing":
+            mesh_efeat_processed, mesh_nfeat_processed = self.processor_encoder(
+                mesh_efeat_embedded,
+                mesh_nfeat_encoded,
+                self.mesh_graph,
+            )
+        else:
+            mesh_nfeat_processed = self.processor_encoder(
+                mesh_nfeat_encoded,
+            )
+            mesh_efeat_processed = None
+        return mesh_efeat_processed, mesh_nfeat_processed, grid_nfeat_encoded
+
+    def decoder_forward(
+        self,
+        mesh_efeat_processed: Tensor,
+        mesh_nfeat_processed: Tensor,
+        grid_nfeat_encoded: Tensor,
+    ) -> Tensor:
+        """Forward method for the last layer of the processor, the decoder,
+        and the final MLP.
+
+        Parameters
+        ----------
+        mesh_efeat_processed : Tensor
+            Multimesh edge features processed by the processor.
+        mesh_nfeat_processed : Tensor
+            Multi-mesh node features processed by the processor.
+        grid_nfeat_encoded : Tensor
+            The encoded node features for the latitude-longitude grid.
+
+        Returns
+        -------
+        grid_nfeat_finale: Tensor
+            The final node features for the latitude-longitude grid.
+        """
+
+        # process multimesh graph
+        if self.processor_type == "MessagePassing":
+            _, mesh_nfeat_processed = self.processor_decoder(
+                mesh_efeat_processed,
+                mesh_nfeat_processed,
+                self.mesh_graph,
+            )
+        else:
+            mesh_nfeat_processed = self.processor_decoder(
+                mesh_nfeat_processed,
+            )
+
+        m2g_efeat_embedded = self.decoder_embedder(self.m2g_edata)
+
+        # decode multimesh to lat/lon
+        grid_nfeat_decoded = self.decoder(
+            m2g_efeat_embedded, grid_nfeat_encoded, mesh_nfeat_processed, self.m2g_graph
+        )
+
+        # map to the target output dimension
+        grid_nfeat_finale = self.finale(
+            grid_nfeat_decoded,
+        )
+
+        return grid_nfeat_finale
+
+    def custom_forward(self, grid_nfeat: Tensor) -> Tensor:
+        """GraphCast forward method with support for gradient checkpointing.
+
+        Parameters
+        ----------
+        grid_nfeat : Tensor
+            Node features of the latitude-longitude graph.
+
+        Returns
+        -------
+        grid_nfeat_finale: Tensor
+            Predicted node features of the latitude-longitude graph.
+        """
+        (
+            mesh_efeat_processed,
+            mesh_nfeat_processed,
+            grid_nfeat_encoded,
+        ) = self.encoder_checkpoint_fn(
+            self.encoder_forward,
+            grid_nfeat,
+            use_reentrant=False,
+            preserve_rng_state=False,
+        )
+
+        # checkpoint of processor done in processor itself
+        if self.processor_type == "MessagePassing":
+            mesh_efeat_processed, mesh_nfeat_processed = self.processor(
+                mesh_efeat_processed,
+                mesh_nfeat_processed,
+                self.mesh_graph,
+            )
+        else:
+            mesh_nfeat_processed = self.processor(
+                mesh_nfeat_processed,
+            )
+            mesh_efeat_processed = None
+
+        grid_nfeat_finale = self.decoder_checkpoint_fn(
+            self.decoder_forward,
+            mesh_efeat_processed,
+            mesh_nfeat_processed,
+            grid_nfeat_encoded,
+            use_reentrant=False,
+            preserve_rng_state=False,
+        )
+
+        return grid_nfeat_finale
+
+    def forward(
+        self,
+        grid_nfeat: Tensor,
+    ) -> Tensor:
+        invar = self.prepare_input(
+            grid_nfeat, self.expect_partitioned_input, self.global_features_on_rank_0
+        )
+        outvar = self.model_checkpoint_fn(
+            self.custom_forward,
+            invar,
+            use_reentrant=False,
+            preserve_rng_state=False,
+        )
+        outvar = self.prepare_output(
+            outvar,
+            self.produce_aggregated_output,
+            self.produce_aggregated_output_on_all_ranks,
+        )
+        return outvar
+
+    def prepare_input(
+        self,
+        invar: Tensor,
+        expect_partitioned_input: bool,
+        global_features_on_rank_0: bool,
+    ) -> Tensor:
+        """Prepares the input to the model in the required shape.
+
+        Parameters
+        ----------
+        invar : Tensor
+            Input in the shape [N, C, H, W].
+
+        expect_partitioned_input : bool
+            flag indicating whether input is partioned according to graph partitioning scheme
+
+        global_features_on_rank_0 : bool
+            Flag indicating whether input is in its "global" form only on group_rank 0 which
+            requires a scatter operation beforehand. Note that only either this flag or
+            expect_partitioned_input can be set at a time.
+
+        Returns
+        -------
+        Tensor
+            Reshaped input.
+        """
+        if global_features_on_rank_0 and expect_partitioned_input:
+            raise ValueError(
+                "global_features_on_rank_0 and expect_partitioned_input cannot be set at the same time."
+            )
+
+        if not self.is_distributed:
+            if invar.size(0) != 1:
+                raise ValueError("GraphCast does not support batch size > 1")
+            invar = invar[0].view(self.input_dim_grid_nodes, -1).permute(1, 0)
+
+        else:
+            # is_distributed
+            if not expect_partitioned_input:
+                # global_features_on_rank_0
+                if invar.size(0) != 1:
+                    raise ValueError("GraphCast does not support batch size > 1")
+
+                invar = invar[0].view(self.input_dim_grid_nodes, -1).permute(1, 0)
+
+                # scatter global features
+                invar = self.g2m_graph.get_src_node_features_in_partition(
+                    invar,
+                    scatter_features=global_features_on_rank_0,
+                )
+
+        return invar
+
+    def prepare_output(
+        self,
+        outvar: Tensor,
+        produce_aggregated_output: bool,
+        produce_aggregated_output_on_all_ranks: bool = True,
+    ) -> Tensor:
+        """Prepares the output of the model in the shape [N, C, H, W].
+
+        Parameters
+        ----------
+        outvar : Tensor
+            Output of the final MLP of the model.
+
+        produce_aggregated_output : bool
+            flag indicating whether output is gathered onto each rank
+            or kept distributed
+
+        produce_aggregated_output_on_all_ranks : bool
+            flag indicating whether output is gatherered on each rank
+            or only gathered at group_rank 0, True by default and
+            only valid if produce_aggregated_output is set.
+
+        Returns
+        -------
+        Tensor
+            The reshaped output of the model.
+        """
+        if produce_aggregated_output or not self.is_distributed:
+            # default case: output of shape [N, C, H, W]
+            if self.is_distributed:
+                outvar = self.m2g_graph.get_global_dst_node_features(
+                    outvar,
+                    get_on_all_ranks=produce_aggregated_output_on_all_ranks,
+                )
+
+            outvar = outvar.permute(1, 0)
+            outvar = outvar.view(self.output_dim_grid_nodes, *self.input_res)
+            outvar = torch.unsqueeze(outvar, dim=0)
+
+        return outvar
+
+    def to(self, *args: Any, **kwargs: Any) -> Self:
+        """Moves the object to the specified device, dtype, or format.
+        This method moves the object and its underlying graph and graph features to
+        the specified device, dtype, or format, and returns the updated object.
+
+        Parameters
+        ----------
+        *args : Any
+            Positional arguments to be passed to the `torch._C._nn._parse_to` function.
+        **kwargs : Any
+            Keyword arguments to be passed to the `torch._C._nn._parse_to` function.
+
+        Returns
+        -------
+        GraphCastNet
+            The updated object after moving to the specified device, dtype, or format.
+        """
+        self = super(GraphCastNet, self).to(*args, **kwargs)
+
+        self.g2m_edata = self.g2m_edata.to(*args, **kwargs)
+        self.m2g_edata = self.m2g_edata.to(*args, **kwargs)
+        self.mesh_ndata = self.mesh_ndata.to(*args, **kwargs)
+        self.mesh_edata = self.mesh_edata.to(*args, **kwargs)
+
+        device, _, _, _ = torch._C._nn._parse_to(*args, **kwargs)
+        self.g2m_graph = self.g2m_graph.to(device)
+        self.mesh_graph = self.mesh_graph.to(device)
+        self.m2g_graph = self.m2g_graph.to(device)
+
+        return self

model/graph_cast_processor.py

@@ -0,0 +1,238 @@
+from typing import Union
+
+import torch
+import torch.nn as nn
+# import transformer_engine as te
+from dgl import DGLGraph
+from torch import Tensor
+
+from onescience.modules.edge.mesh_edge_block import MeshEdgeBlock
+from onescience.modules.node.mesh_node_block import MeshNodeBlock
+from onescience.modules.utils.gnnlayer_utils import CuGraphCSC, set_checkpoint_fn
+
+
+class GraphCastProcessor(nn.Module):
+    """Processor block used in GraphCast operating on a latent space
+    represented by hierarchy of icosahedral meshes.
+
+    Parameters
+    ----------
+    aggregation : str, optional
+        message passing aggregation method ("sum", "mean"), by default "sum"
+    processor_layers : int, optional
+        number of processor layers, by default 16
+    input_dim_nodes : int, optional
+        input dimensionality of the node features, by default 512
+    input_dim_edges : int, optional
+        input dimensionality of the edge features, by default 512
+    hidden_dim : int, optional
+        number of neurons in each hidden layer, by default 512
+    hidden_layers : int, optional
+        number of hiddel layers, by default 1
+    activation_fn : nn.Module, optional
+        type of activation function, by default nn.SiLU()
+    norm_type : str, optional
+        Normalization type ["TELayerNorm", "LayerNorm"].
+        Use "TELayerNorm" for optimal performance. By default "LayerNorm".
+    do_conat_trick: : bool, default=False
+        whether to replace concat+MLP with MLP+idx+sum
+    recompute_activation : bool, optional
+        Flag for recomputing activation in backward to save memory, by default False.
+        Currently, only SiLU is supported.
+    """
+
+    def __init__(
+        self,
+        aggregation: str = "sum",
+        processor_layers: int = 16,
+        input_dim_nodes: int = 512,
+        input_dim_edges: int = 512,
+        hidden_dim: int = 512,
+        hidden_layers: int = 1,
+        activation_fn: nn.Module = nn.SiLU(),
+        norm_type: str = "LayerNorm",
+        do_concat_trick: bool = False,
+        recompute_activation: bool = False,
+    ):
+        super().__init__()
+
+        edge_block_invars = (
+            input_dim_nodes,
+            input_dim_edges,
+            input_dim_edges,
+            hidden_dim,
+            hidden_layers,
+            activation_fn,
+            norm_type,
+            do_concat_trick,
+            recompute_activation,
+        )
+        node_block_invars = (
+            aggregation,
+            input_dim_nodes,
+            input_dim_edges,
+            input_dim_nodes,
+            hidden_dim,
+            hidden_layers,
+            activation_fn,
+            norm_type,
+            recompute_activation,
+        )
+
+        layers = []
+        for _ in range(processor_layers):
+            layers.append(MeshEdgeBlock(**dict(zip(
+              ['input_dim_nodes','input_dim_edges','output_dim','hidden_dim',
+               'hidden_layers','activation_fn','norm_type','do_concat_trick',
+               'recompute_activation'], edge_block_invars))))
+            layers.append(MeshNodeBlock(**dict(zip(
+              ['aggregation','input_dim_nodes','input_dim_edges','output_dim',
+               'hidden_dim','hidden_layers','activation_fn','norm_type',
+               'recompute_activation'], node_block_invars))))
+
+        self.processor_layers = nn.ModuleList(layers)
+        self.num_processor_layers = len(self.processor_layers)
+        # per default, no checkpointing
+        # one segment for compatability
+        self.checkpoint_segments = [(0, self.num_processor_layers)]
+        self.checkpoint_fn = set_checkpoint_fn(False)
+
+    def set_checkpoint_segments(self, checkpoint_segments: int):
+        """
+        Set the number of checkpoint segments
+
+        Parameters
+        ----------
+        checkpoint_segments : int
+            number of checkpoint segments
+
+        Raises
+        ------
+        ValueError
+            if the number of processor layers is not a multiple of the number of
+            checkpoint segments
+        """
+        if checkpoint_segments > 0:
+            if self.num_processor_layers % checkpoint_segments != 0:
+                raise ValueError(
+                    "Processor layers must be a multiple of checkpoint_segments"
+                )
+            segment_size = self.num_processor_layers // checkpoint_segments
+            self.checkpoint_segments = []
+            for i in range(0, self.num_processor_layers, segment_size):
+                self.checkpoint_segments.append((i, i + segment_size))
+
+            self.checkpoint_fn = set_checkpoint_fn(True)
+        else:
+            self.checkpoint_fn = set_checkpoint_fn(False)
+            self.checkpoint_segments = [(0, self.num_processor_layers)]
+
+    def run_function(self, segment_start: int, segment_end: int):
+        """Custom forward for gradient checkpointing
+
+        Parameters
+        ----------
+        segment_start : int
+            Layer index as start of the segment
+        segment_end : int
+            Layer index as end of the segment
+
+        Returns
+        -------
+        function
+            Custom forward function
+        """
+        segment = self.processor_layers[segment_start:segment_end]
+
+        def custom_forward(efeat, nfeat, graph):
+            """Custom forward function"""
+            for module in segment:
+                efeat, nfeat = module(efeat, nfeat, graph)
+            return efeat, nfeat
+
+        return custom_forward
+
+    def forward(
+        self,
+        efeat: Tensor,
+        nfeat: Tensor,
+        graph: Union[DGLGraph, CuGraphCSC],
+    ) -> Tensor:
+        for segment_start, segment_end in self.checkpoint_segments:
+            efeat, nfeat = self.checkpoint_fn(
+                self.run_function(segment_start, segment_end),
+                efeat,
+                nfeat,
+                graph,
+                use_reentrant=False,
+                preserve_rng_state=False,
+            )
+
+        return efeat, nfeat
+
+
+class GraphCastProcessorGraphTransformer(nn.Module):
+    """Processor block used in GenCast operating on a latent space
+    represented by hierarchy of icosahedral meshes.
+
+    Parameters
+    ----------
+    attn_mask : torch.Tensor
+        Attention mask to be applied within the transformer layers.
+    processor_layers : int, optional (default=16)
+        Number of processing layers.
+    input_dim_nodes : int, optional (default=512)
+        Dimension of the input features for each node.
+    hidden_dim : int, optional (default=512)
+        Dimension of the hidden features within the transformer layers.
+    """
+
+    def __init__(
+        self,
+        attention_mask: torch.Tensor,
+        num_attention_heads: int = 4,
+        processor_layers: int = 16,
+        input_dim_nodes: int = 512,
+        hidden_dim: int = 512,
+    ):
+        super().__init__()
+        self.num_attention_heads = num_attention_heads
+        self.hidden_dim = hidden_dim
+        self.attention_mask = torch.tensor(attention_mask, dtype=torch.bool)
+        self.register_buffer("mask", self.attention_mask, persistent=False)
+
+        layers = [
+        #     te.pytorch.TransformerLayer(
+        #         hidden_size=input_dim_nodes,
+        #         ffn_hidden_size=hidden_dim,
+        #         num_attention_heads=num_attention_heads,
+        #         layer_number=i + 1,
+        #         fuse_qkv_params=False,
+        #     )
+        #     for i in range(processor_layers)
+
+            nn.TransformerEncoderLayer(
+                d_model=input_dim_nodes,           # 等同于 hidden_size
+                nhead=num_attention_heads,         # 等同于 num_attention_heads
+                dim_feedforward=hidden_dim,        # 等同于 ffn_hidden_size
+                activation='gelu',                 # 激活函数可以选择 gelu 或 relu
+                batch_first=True                   # 使得输入张量维度为 (batch_size, seq_len, d_model)
+            )
+            for _ in range(processor_layers)       # 创建多层 Transformer 编码器层
+        ]
+        self.processor_layers = nn.ModuleList(layers)
+
+    def forward(
+        self,
+        nfeat: Tensor,
+    ) -> Tensor:
+        nfeat = nfeat.unsqueeze(1)
+        # TODO make sure reshaping the last dim to (h, d) is done automatically in the transformer layer
+        for module in self.processor_layers:
+            nfeat = module(
+                nfeat,
+                attention_mask=self.mask,
+                self_attn_mask_type="arbitrary",
+            )
+
+        return torch.squeeze(nfeat, 1)

scripts/compute_time_diff_std.py

@@ -0,0 +1,48 @@
+import torch
+import os
+import sys
+import numpy as np
+from tqdm import tqdm
+from onescience.datapipes.climate import ERA5Datapipe
+from onescience.utils.YParams import YParams
+
+
+def main():
+    # instantiate the training datapipe
+    config_file_path = os.path.join(current_path, 'conf/config.yaml')
+    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=False
+    )
+    train_dataloader, train_sampler = datapipe.get_dataloader("train")
+    
+    print(f"Loaded training datapipe of length {len(train_dataloader)}")
+
+    area = torch.abs(torch.cos(torch.linspace(-90, 90, steps=cfg_data.dataset.img_size[0]) * np.pi / 180))
+    area /= torch.mean(area)
+    area = area.unsqueeze(1)
+
+    mean, mean_sqr = 0, 0
+    for data in tqdm(train_dataloader):
+        invar = data[0]  # [b, N, h, w]
+        outvar = data[1]  # [b, N, h, w]
+        diff = outvar - invar
+        weighted_diff = area * diff
+        weighted_diff_sqr = torch.square(weighted_diff)
+        mean += torch.mean(weighted_diff, dim=(2, 3)) / len(train_dataloader)
+        mean_sqr += torch.mean(weighted_diff_sqr, dim=(2, 3)) / len(train_dataloader)
+
+    variance = mean_sqr - mean**2  # [1,num_channel, 1,1]
+    std = torch.sqrt(variance)
+
+    np.save("time_diff_std.npy", std.numpy())
+    print(f"saving time_diff_std.npy, shapes are {std.numpy().shape}")
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    main()

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/finetune.py

@@ -0,0 +1,331 @@
+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 numpy as np
+import torch.distributed as dist
+import logging
+import time
+
+from torch.nn.parallel import DistributedDataParallel
+from torch.optim.lr_scheduler import SequentialLR, LinearLR, CosineAnnealingLR, LambdaLR
+from onescience.datapipes.climate import ERA5Datapipe
+from onescience.utils.YParams import YParams
+from onescience.modules.utils.graphcast.data_utils import StaticData
+from onescience.modules.utils.graphcast.graph_utils import deg2rad
+from model.graph_cast_net import GraphCastNet
+from onescience.modules.utils.graphcast.loss import GraphCastLossFunction
+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")
+
+    input_dim_grid_nodes = (len(cfg_data.dataset.channels) + cfg.use_cos_zenith + 4 * cfg.use_time_of_year_index) * (cfg.num_history + 1) + cfg.num_channels_static
+    
+    model = GraphCastNet(mesh_level=cfg.mesh_level,
+                         multimesh=cfg.multimesh,
+                         input_res=tuple(cfg_data.dataset.img_size),
+                         input_dim_grid_nodes=input_dim_grid_nodes,
+                         input_dim_mesh_nodes=3,
+                         input_dim_edges=4,
+                         output_dim_grid_nodes=len(cfg_data.dataset.channels),
+                         processor_type=cfg.processor_type,
+                         khop_neighbors=cfg.khop_neighbors,
+                         num_attention_heads=cfg.num_attention_heads,
+                         processor_layers=cfg.processor_layers,
+                         hidden_dim=cfg.hidden_dim,
+                         norm_type=cfg.norm_type,
+                         do_concat_trick=cfg.concat_trick,
+                         recompute_activation=cfg.recompute_activation,
+                         )
+    model_dtype = torch.bfloat16 if cfg.full_bf16 else torch.float32
+    model.set_checkpoint_encoder(cfg.checkpoint_encoder)
+    model.set_checkpoint_decoder(cfg.checkpoint_decoder)
+    model = model.to(dtype=model_dtype).to(local_rank)
+    if hasattr(model, "module"):
+        latitudes = model.module.latitudes
+        longitudes = model.module.longitudes
+        lat_lon_grid = model.module.lat_lon_grid
+    else:
+        latitudes = model.latitudes
+        longitudes = model.longitudes
+        lat_lon_grid = model.lat_lon_grid
+    tatic_dir = os.path.join(cfg_data.dataset.data_dir, "static")
+    
+    static_data = StaticData(static_dir, latitudes, longitudes).get().to(device=local_rank)
+    channels_list = [i for i in range(len(cfg_data.dataset.channels))]
+    area = torch.abs(torch.cos(deg2rad(lat_lon_grid[:, :, 0])))
+    area /= torch.mean(area)
+    area = area.to(dtype=torch.bfloat16 if cfg.full_bf16 else torch.float32).to(device=local_rank)
+    criterion = GraphCastLossFunction(area, channels_list, cfg_data.dataset.dataset_metadata_path, cfg_data.dataset.time_diff_std_path)
+    optimizer = optimizers.FusedAdam(model.parameters(),
+                                     lr=cfg.lr, betas=(0.9, 0.95),
+                                     adam_w_mode=True,
+                                     weight_decay=0.1)
+    scheduler1 = LinearLR(optimizer, start_factor=1e-3, end_factor=1.0, total_iters=cfg.num_iters_step1, )
+    scheduler2 = CosineAnnealingLR(optimizer, T_max=cfg.num_iters_step2, eta_min=0.0)
+    scheduler3 = LambdaLR(optimizer, lr_lambda=lambda epoch: (cfg.lr_step3 / cfg.lr))
+    scheduler = SequentialLR(optimizer,
+                             schedulers=[scheduler1, scheduler2, scheduler3],
+                             milestones=[cfg.num_iters_step1, cfg.num_iters_step1 + cfg.num_iters_step2])
+    
+    ## Train process init
+    os.makedirs(cfg.checkpoint_dir, exist_ok=True)
+    train_loss_file = f"{cfg.checkpoint_dir}/ft_trloss.npy"
+    best_valid_loss = 1.0e6
+    best_loss_epoch = 0
+    train_losses = np.empty((0,), dtype=np.float32)
+
+    ## Get model params count
+    if cfg.world_size == 1:
+        total_params = sum(p.numel() for p in model.parameters())
+        print("\n\n")
+        print("-" * 50)
+        print(f"📂 now params is {total_params}, {total_params / 1e6:.2f}M, {total_params / 1e9:.2f}B")
+        print("-" * 50, "\n")
+
+    ## Load model weight if there exist well-trained model 
+    if not os.path.exists(f"{cfg.checkpoint_dir}/model_bak.pth"):
+        print('⚠️ ⚠️ Please train the model first...')
+        exit()
+    
+    if os.path.exists(f"{cfg.checkpoint_dir}/model_finetune_bak.pth"):
+        if world_rank == 0:
+            print("\n\n")
+            print("-" * 50)
+            print(f"✅ There has a finetune-model weight, load and continue training...")
+            print(f'If you want to finetune a new model, ensure there is no model_finetune_bak.pth file in {cfg.checkpoint_dir}')
+            print("-" * 50, "\n")
+        ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location=f'cuda:{local_rank}', weights_only=False)
+        model.load_state_dict(ckpt["model_state_dict"])
+        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+        scheduler.load_state_dict(ckpt["scheduler_state_dict"])
+        best_valid_loss = ckpt["best_valid_loss"]
+        best_loss_epoch = ckpt["best_loss_epoch"]
+        train_losses = np.load(train_loss_file)
+    else:
+        if world_rank == 0:
+            print("\n\n")
+            print("-" * 50)
+            print(f"✅ Load well-trained model and continue to finetune...")
+            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"])
+
+    if cfg.world_size > 1:
+        model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
+    if cfg.force_single_checkpoint_finetune:
+        if hasattr(model, "module"):
+            model.module.set_checkpoint_model(True)
+        else:
+            model.set_checkpoint_model(True)
+    if cfg.checkpoint_encoder_finetune:
+        if hasattr(model, "module"):
+            model.module.set_checkpoint_encoder(True)
+        else:
+            model.set_checkpoint_encoder(True)
+    if cfg.checkpoint_processor_finetune:
+        if hasattr(model, "module"):
+            model.module.set_checkpoint_processor(cfg.segments)
+        else:
+            model.set_checkpoint_encoder(True)
+    if cfg.checkpoint_decoder_finetune:
+        if hasattr(model, "module"):
+            model.module.set_checkpoint_decoder(True)
+        else:
+            model.set_checkpoint_encoder(True)
+
+    world_rank == 0 and logger.info(f"start finetuning ...")
+
+
+    num_rollout_steps = 2
+    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,
+        distributed=dist.is_initialized(),
+        output_steps=num_rollout_steps,
+        num_workers=0
+    )
+    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,
+        num_workers=0
+    )
+    val_dataloader, val_sampler = datapipe.get_dataloader("valid")
+
+    for epoch in range(cfg.num_iters_step3):
+        if epoch % cfg.step_change_freq == 0:
+            num_rollout_steps = epoch // cfg.step_change_freq + 2
+            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,
+                distributed=dist.is_initialized(),
+                output_steps=num_rollout_steps,
+                num_workers=0
+            )
+            train_dataloader, train_sampler = datapipe.get_dataloader("train")
+
+        if dist.is_initialized():
+            train_sampler.set_epoch(epoch)
+            val_sampler.set_epoch(epoch)
+        train_loss = 0
+        start_time = time.time()
+        for j, data in enumerate(train_dataloader):
+            invar = data[0].to(device=local_rank)
+            
+            outvar = data[1].to(device=local_rank)
+            cos_zenith = data[2].to(device=local_rank)
+            in_idx = data[3].item()
+            cos_zenith = torch.squeeze(cos_zenith, dim=2)
+            cos_zenith = torch.clamp(cos_zenith, min=0.0) - 1.0 / torch.pi
+            outvar = outvar.to(dtype=model_dtype)
+            for t in range(outvar.shape[1]):
+                day_of_year, time_of_day = divmod(in_idx + t * cfg.dt, 24 // cfg.dt)
+                normalized_day_of_year = torch.tensor((day_of_year / 365) * (np.pi / 2), dtype=torch.float32, device=local_rank)
+                normalized_time_of_day = torch.tensor((time_of_day / (24 - cfg.dt)) * (np.pi / 2), dtype=torch.float32, device=local_rank)
+                sin_day_of_year = torch.sin(normalized_day_of_year).expand(1, 1, 721, 1440)
+                cos_day_of_year = torch.cos(normalized_day_of_year).expand(1, 1, 721, 1440)
+                sin_time_of_day = torch.sin(normalized_time_of_day).expand(1, 1, 721, 1440)
+                cos_time_of_day = torch.cos(normalized_time_of_day).expand(1, 1, 721, 1440)
+                invar = torch.concat((invar, cos_zenith[:, t:t + 1, :, :], static_data, sin_day_of_year, cos_day_of_year, sin_time_of_day, cos_time_of_day), dim=1)
+                invar = invar.to(dtype=model_dtype)
+                if t < outvar.shape[1] - 1:
+                    with torch.no_grad():
+                        outpred = model(invar)
+                    invar = outpred.detach()
+                else:
+                    outpred = model(invar)
+                    loss = criterion(outpred, outvar[:, t])
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip_norm)
+            torch.cuda.nvtx.range_pop()
+            optimizer.step()
+            scheduler.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}')
+
+            if (j + 1) % cfg.val_freq == 0:
+                model.eval()
+                valid_loss = 0.0
+                with torch.no_grad():
+                    start_time = time.time()
+                    for k, data in enumerate(val_dataloader):
+                        invar = data[0].to(device=local_rank)
+                        outvar = data[1].to(device=local_rank)
+                        cos_zenith = data[2].to(device=local_rank)
+                        in_idx = data[3].item()
+
+                        cos_zenith = torch.squeeze(cos_zenith, dim=2)
+                        cos_zenith = torch.clamp(cos_zenith, min=0.0) - 1.0 / torch.pi  # [b, 2, h, w]
+                        outvar = outvar.to(dtype=model_dtype)
+                        loss = 0.0
+                        for t in range(outvar.shape[1]):
+                            day_of_year, time_of_day = divmod(in_idx + t * cfg.dt, 24 // cfg.dt)
+                            normalized_day_of_year = torch.tensor((day_of_year / 365) * (np.pi / 2), dtype=torch.float32, device=local_rank)
+                            normalized_time_of_day = torch.tensor((time_of_day / (24 - cfg.dt)) * (np.pi / 2), dtype=torch.float32, device=local_rank)
+                            sin_day_of_year = torch.sin(normalized_day_of_year).expand(1, 1, 721, 1440)
+                            cos_day_of_year = torch.cos(normalized_day_of_year).expand(1, 1, 721, 1440)
+                            sin_time_of_day = torch.sin(normalized_time_of_day).expand(1, 1, 721, 1440)
+                            cos_time_of_day = torch.cos(normalized_time_of_day).expand(1, 1, 721, 1440)
+                            invar = torch.concat((invar, cos_zenith[:, t:t + 1, :, :], static_data, sin_day_of_year, cos_day_of_year, sin_time_of_day, cos_time_of_day), dim=1)
+                            invar = invar.to(dtype=model_dtype)
+                            outpred = model(invar)
+                            invar = outpred
+                            loss += criterion(outpred, outvar[:, t])
+
+                        loss /= outvar.shape[1]
+                        if cfg.world_size > 1:
+                            loss_tensor = loss.detach().to(local_rank)  # torch.tensor(loss, device=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} {k+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)/(k+1): .02f}s/{cfg_data.dataloader.batch_size}batch] '
+                                    f'loss:{valid_loss / (k+1): .04f}')
+                        
+                    valid_loss /= len(val_dataloader)
+                    is_save_ckp = False
+                    if valid_loss < best_valid_loss:
+                        best_valid_loss = valid_loss
+                        best_loss_epoch = epoch
+                        world_rank == 0 and save_checkpoint(model, optimizer, scheduler, best_valid_loss, best_loss_epoch, cfg.checkpoint_dir)
+                        is_save_ckp = True
+                train_loss /= (j+1)
+                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)
+                    np.save(train_loss_file, train_losses)
+                
+
+    print(f'Graphcast has been well-trained, next step is fine-tune')
+
+
+def save_checkpoint(model, optimizer, scheduler, best_valid_loss, best_loss_epoch, model_path):
+    model_to_save = model.module if hasattr(model, "module") else model
+    state = {"model_state_dict": model_to_save.state_dict(),
+             "optimizer_state_dict": optimizer.state_dict(),
+             "scheduler_state_dict": scheduler.state_dict(),
+             "best_valid_loss": best_valid_loss,
+             "best_loss_epoch": best_loss_epoch,
+            }
+    torch.save(state, f"{model_path}/model_finetune.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_finetune.pth {model_path}/model_finetune_bak.pth")
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    main()

scripts/get_data_json.py

@@ -0,0 +1,24 @@
+import os
+from onescience.utils.YParams import YParams
+import json
+import sys
+
+
+def main():
+    config_file_path = os.path.join(current_path, 'conf/config.yaml')
+    cfg_data = YParams(config_file_path, "datapipe")
+    metadata = {
+        "coords": {
+            "channel": {str(i): name for i, name in enumerate(cfg_data.dataset.channels)}
+        }
+    }
+    with open('./data.json', "w") as f:
+        json.dump(metadata, f, indent=4)
+    print(f"Metadata saved.")
+
+
+
+if __name__ == '__main__':
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    main()

scripts/inference.py

@@ -0,0 +1,122 @@
+import sys
+from pathlib import Path
+
+# 获取项目根目录（train.py上级的上级）
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+import os
+import numpy as np
+import glob
+import torch
+import sys
+import h5py
+from tqdm import tqdm
+
+from onescience.utils.YParams import YParams
+from onescience.datapipes.climate import ERA5Datapipe
+from ruamel.yaml.scalarfloat import ScalarFloat
+from onescience.modules.utils.graphcast.data_utils import StaticData
+from onescience.modules.utils.graphcast.graph_utils import deg2rad
+from model.graph_cast_net import GraphCastNet
+
+
+torch.serialization.add_safe_globals([ScalarFloat])
+
+
+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")
+
+    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")
+    means, stds = get_stats(cfg_data.dataset.data_dir, cfg_data.dataset.channels)
+
+    ckpt = torch.load(f"{cfg.checkpoint_dir}/model_finetune_bak.pth", map_location="cuda:0", weights_only=True)
+    model_dtype = torch.bfloat16 if cfg.full_bf16 else torch.float32
+    input_dim_grid_nodes = (len(cfg_data.dataset.channels) + cfg.use_cos_zenith + 4 * cfg.use_time_of_year_index) * (cfg.num_history + 1) + cfg.num_channels_static
+    model = GraphCastNet(mesh_level=cfg.mesh_level,
+                         multimesh=cfg.multimesh,
+                         input_res=tuple(cfg_data.dataset.img_size),
+                         input_dim_grid_nodes=input_dim_grid_nodes,
+                         input_dim_mesh_nodes=3,
+                         input_dim_edges=4,
+                         output_dim_grid_nodes=len(cfg_data.dataset.channels),
+                         processor_type=cfg.processor_type,
+                         khop_neighbors=cfg.khop_neighbors,
+                         num_attention_heads=cfg.num_attention_heads,
+                         processor_layers=cfg.processor_layers,
+                         hidden_dim=cfg.hidden_dim,
+                         norm_type=cfg.norm_type,
+                         do_concat_trick=cfg.concat_trick,
+                         recompute_activation=cfg.recompute_activation,
+                         )
+
+    model.set_checkpoint_encoder(cfg.checkpoint_encoder)
+    model.set_checkpoint_decoder(cfg.checkpoint_decoder)
+    model = model.to(dtype=model_dtype).to("cuda:0")
+    model.load_state_dict(ckpt["model_state_dict"])
+
+    if hasattr(model, "module"):
+        latitudes = model.module.latitudes
+        longitudes = model.module.longitudes
+    else:
+        latitudes = model.latitudes
+        longitudes = model.longitudes
+    static_dir = os.path.join(cfg_data.dataset.data_dir, "static")
+    
+    static_data = StaticData(static_dir, latitudes, longitudes).get().to(device=local_rank)
+    model.eval()
+    os.makedirs('./result/output/', exist_ok=True)
+    print(f"📂 samples will be generated to './result/output/'")
+    with torch.no_grad():
+        for data in tqdm(test_dataloader, desc="Inferring testset", unit="batch"):
+            invar = data[0].to(device="cuda:0")
+            cos_zenith = data[2].to(device="cuda:0")
+            in_idx = data[3].item()
+            filename = data[4][-1][0]
+
+            cos_zenith = torch.squeeze(cos_zenith, dim=2)
+            cos_zenith = torch.clamp(cos_zenith, min=0.0) - 1.0 / torch.pi
+            day_of_year, time_of_day = divmod(in_idx * cfg.dt, 24)
+            normalized_day_of_year = torch.tensor((day_of_year / 365) * (np.pi / 2), dtype=torch.float32, device="cuda:0")
+            normalized_time_of_day = torch.tensor((time_of_day / (24 - cfg.dt)) * (np.pi / 2), dtype=torch.float32, device="cuda:0")
+            sin_day_of_year = torch.sin(normalized_day_of_year).expand(1, 1, 721, 1440)
+            cos_day_of_year = torch.cos(normalized_day_of_year).expand(1, 1, 721, 1440)
+            sin_time_of_day = torch.sin(normalized_time_of_day).expand(1, 1, 721, 1440)
+            cos_time_of_day = torch.cos(normalized_time_of_day).expand(1, 1, 721, 1440)
+            invar = torch.concat((invar, cos_zenith, static_data, sin_day_of_year, cos_day_of_year, sin_time_of_day, cos_time_of_day), dim=1)
+
+            invar = invar.to(dtype=model_dtype)
+            pred_var = model(invar).to(dtype=torch.float32)
+            pred_var = pred_var.cpu().numpy()
+            pred_var = pred_var * stds + means
+
+            np.save(f"result/output/{filename}.npy", pred_var)

scripts/result.py

@@ -0,0 +1,196 @@
+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 get_result(total_files, channel_indices, time_step, data_dir, clim_mean):
+    channel_rmse = np.zeros(len(channel_indices))
+    channel_acc = np.zeros(len(channel_indices))
+    clim_mean = clim_mean[0, :, :, :]
+    if not os.path.exists('result/rmse.npy') or not os.path.exists('result/acc.npy'):
+        numerator = np.zeros(len(channel_indices))
+        pred_sq_sum = np.zeros(len(channel_indices))
+        label_sq_sum = np.zeros(len(channel_indices))
+        for file in tqdm(total_files, unit="files"):
+            fname = file[:-4]  # 去掉 .npy
+            year = fname[:4]
+            t_idx = filename_to_index(fname, time_step)
+            with h5py.File(os.path.join(data_dir, 'data', f'{year}.h5'), "r") as f:
+                label = f["fields"][t_idx]  # [C, H, W]
+                label = label[channel_indices]
+            pred = np.load(f'result/output/{file}').squeeze()
+
+            label_anom = label - clim_mean
+            pred_anom = pred - clim_mean
+            # 累加
+            numerator += np.sum(pred_anom * label_anom, axis=(1, 2))
+            pred_sq_sum += np.sum(pred_anom ** 2, axis=(1, 2))
+            label_sq_sum += np.sum(label_anom ** 2, axis=(1, 2))
+
+            channel_rmse += np.sqrt(np.mean((label - pred) ** 2, axis=(1, 2)))
+        channel_rmse /= len(total_files)
+        channel_acc = numerator / (np.sqrt(pred_sq_sum * label_sq_sum) + 1e-8)
+        np.save('./result/acc.npy', channel_acc)
+        np.save('result/rmse.npy', channel_rmse)
+
+
+def show_result():
+    channel_rmse = np.load('result/rmse.npy')
+    channel_acc = np.load('./result/acc.npy')
+
+    channels = [cfg_data.dataset.channels[i] for i in range(len(channel_indices))]
+    w = 24  # 最长 channel 名宽度
+
+    # 表头
+    print(f"┌{'─' * (w + 2)}┬{'─' * 14}┬{'─' * 14}┐")
+    print(f"│ {'Channel':<{w}} │ {'RMSE':>12} │ {'ACC':>12} │")
+    print(f"├{'─' * (w + 2)}┼{'─' * 14}┼{'─' * 14}┤")
+    # 数据行
+    for i, ch in enumerate(channels):
+        print(f"│ {ch:<{w}} │ {channel_rmse[i]:>12.4f} | {channel_acc[i]:>12.4f} |")
+    print(f"├{'─' * (w + 2)}┼{'─' * 14}┼{'─' * 14}┤")
+    print(f"│ {'Average':<{w}} │ {np.mean(channel_rmse):>12.4f} │ {np.mean(channel_acc):>12.4f} │")
+    print(f"└{'─' * (w + 2)}┴{'─' * 14}┴{'─' * 14}┘")
+
+
+def plot(label, pred, var, filename):
+    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
+
+    xtick_labels = ['180°W', '90°W', '0°', '90°E', '180°E']
+    ytick_labels = ['90°S', '45°S', '0°', '45°N', '90°N']
+    xticks = np.linspace(0, label.shape[-1] - 1, 5)
+    yticks = np.linspace(0, label.shape[-2] - 1, 5)
+
+    vmin = min(label.min(), pred.min())
+    vmax = max(label.max(), pred.max())
+
+    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):
+    mask = ~(np.isnan(train_loss))
+    train_loss = train_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')
+    min_idx = np.argmin(train_loss)
+    ax.scatter(epochs[min_idx], train_loss[min_idx],
+               color=colors['valid'], s=40, zorder=5, edgecolors='white')
+    ax.annotate(f'Best: {train_loss[min_idx]:.3f}',
+                xy=(epochs[min_idx], train_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')
+    plot_loss(train_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,276 @@
+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 numpy as np
+import torch.distributed as dist
+import logging
+import time
+
+from torch.nn.parallel import DistributedDataParallel
+from torch.optim.lr_scheduler import SequentialLR, LinearLR, CosineAnnealingLR, LambdaLR
+from onescience.datapipes.climate import ERA5Datapipe
+from onescience.utils.YParams import YParams
+from onescience.modules.utils.graphcast.data_utils import StaticData
+from onescience.modules.utils.graphcast.graph_utils import deg2rad
+from model.graph_cast_net import GraphCastNet
+from onescience.modules.utils.graphcast.loss import GraphCastLossFunction
+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(),
+        num_workers=0
+    )
+    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(),
+        num_workers=0
+    )
+    val_dataloader, val_sampler = datapipe.get_dataloader("valid")
+
+    input_dim_grid_nodes = (len(cfg_data.dataset.channels) + cfg.use_cos_zenith + 4 * cfg.use_time_of_year_index) * (cfg.num_history + 1) + cfg.num_channels_static
+    model = GraphCastNet(mesh_level=cfg.mesh_level,
+                         multimesh=cfg.multimesh,
+                         input_res=tuple(cfg_data.dataset.img_size),
+                         input_dim_grid_nodes=input_dim_grid_nodes,
+                         input_dim_mesh_nodes=3,
+                         input_dim_edges=4,
+                         output_dim_grid_nodes=len(cfg_data.dataset.channels),
+                         processor_type=cfg.processor_type,
+                         khop_neighbors=cfg.khop_neighbors,
+                         num_attention_heads=cfg.num_attention_heads,
+                         processor_layers=cfg.processor_layers,
+                         hidden_dim=cfg.hidden_dim,
+                         norm_type=cfg.norm_type,
+                         do_concat_trick=cfg.concat_trick,
+                         recompute_activation=cfg.recompute_activation,
+                         )
+    model_dtype = torch.bfloat16 if cfg.full_bf16 else torch.float32
+    model.set_checkpoint_encoder(cfg.checkpoint_encoder)
+    model.set_checkpoint_decoder(cfg.checkpoint_decoder)
+    model = model.to(dtype=model_dtype).to(local_rank)
+    if hasattr(model, "module"):
+        latitudes = model.module.latitudes
+        longitudes = model.module.longitudes
+        lat_lon_grid = model.module.lat_lon_grid
+    else:
+        latitudes = model.latitudes
+        longitudes = model.longitudes
+        lat_lon_grid = model.lat_lon_grid
+    
+    static_dir = os.path.join(cfg_data.dataset.data_dir, "static")
+    
+    static_data = StaticData(static_dir, latitudes, longitudes).get().to(device=local_rank)
+    channels_list = [i for i in range(len(cfg_data.dataset.channels))]
+    area = torch.abs(torch.cos(deg2rad(lat_lon_grid[:, :, 0])))
+    area /= torch.mean(area)
+    area = area.to(dtype=torch.bfloat16 if cfg.full_bf16 else torch.float32).to(device=local_rank)
+    criterion = GraphCastLossFunction(area, channels_list, cfg_data.dataset.dataset_metadata_path, cfg_data.dataset.time_diff_std_path)
+    optimizer = optimizers.FusedAdam(model.parameters(),
+                                     lr=cfg.lr, betas=(0.9, 0.95),
+                                     adam_w_mode=True,
+                                     weight_decay=0.1)
+    scheduler1 = LinearLR(optimizer, start_factor=1e-3, end_factor=1.0, total_iters=cfg.num_iters_step1, )
+    scheduler2 = CosineAnnealingLR(optimizer, T_max=cfg.num_iters_step2, eta_min=0.0)
+    scheduler3 = LambdaLR(optimizer, lr_lambda=lambda epoch: (cfg.lr_step3 / cfg.lr))
+    scheduler = SequentialLR(optimizer,
+                             schedulers=[scheduler1, scheduler2, scheduler3],
+                             milestones=[cfg.num_iters_step1, cfg.num_iters_step1 + cfg.num_iters_step2])
+    
+    ## Train process init
+    os.makedirs(cfg.checkpoint_dir, exist_ok=True)
+    train_loss_file = f"{cfg.checkpoint_dir}/trloss.npy"
+    best_valid_loss = 1.0e6
+    best_loss_epoch = 0
+    train_losses = np.empty((0,), dtype=np.float32)
+
+    ## Get model params count
+    if cfg.world_size == 1:
+        total_params = sum(p.numel() for p in model.parameters())
+        print("\n\n")
+        print("-" * 50)
+        print(f"📂 now params is {total_params}, {total_params / 1e6:.2f}M, {total_params / 1e9:.2f}B")
+        print("-" * 50, "\n")
+
+    ## Load model weight if there exist well-trained model 
+    if os.path.exists(f"{cfg.checkpoint_dir}/model_bak.pth"):
+        if world_rank == 0:
+            print("\n\n")
+            print("-" * 50)
+            print(f"✅ There has a model weight, load and continue training...")
+            print(f'If you want to train a new model, ensure there is no *.pth file in {cfg.checkpoint_dir}')
+            print("-" * 50, "\n")
+        ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location=f'cuda:{local_rank}', weights_only=False)
+        model.load_state_dict(ckpt["model_state_dict"])
+        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+        scheduler.load_state_dict(ckpt["scheduler_state_dict"])
+        best_valid_loss = ckpt["best_valid_loss"]
+        best_loss_epoch = ckpt["best_loss_epoch"]
+        train_losses = np.load(train_loss_file)
+
+    ## 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(cfg.max_epoch):
+        if dist.is_initialized():
+            train_sampler.set_epoch(epoch)
+            val_sampler.set_epoch(epoch)
+
+        model.train()
+        train_loss = 0
+        start_time = time.time()
+        for j, data in enumerate(train_dataloader):
+            invar = data[0].to(device=local_rank)
+            outvar = data[1].to(device=local_rank)
+            cos_zenith = data[2].to(device=local_rank)
+            in_idx = data[3].item()
+            cos_zenith = torch.squeeze(cos_zenith, dim=2)
+            cos_zenith = torch.clamp(cos_zenith, min=0.0) - 1.0 / torch.pi
+            day_of_year, time_of_day = divmod(in_idx * cfg.dt, 24)
+            normalized_day_of_year = torch.tensor((day_of_year / 365) * (np.pi / 2), dtype=torch.float32, device=local_rank)
+            normalized_time_of_day = torch.tensor((time_of_day / (24 - cfg.dt)) * (np.pi / 2), dtype=torch.float32, device=local_rank)
+            sin_day_of_year = torch.sin(normalized_day_of_year).expand(1, 1, 721, 1440)
+            cos_day_of_year = torch.cos(normalized_day_of_year).expand(1, 1, 721, 1440)
+            sin_time_of_day = torch.sin(normalized_time_of_day).expand(1, 1, 721, 1440)
+            cos_time_of_day = torch.cos(normalized_time_of_day).expand(1, 1, 721, 1440)
+            invar = torch.concat((invar, cos_zenith, static_data, sin_day_of_year, cos_day_of_year, sin_time_of_day, cos_time_of_day), dim=1)
+            invar, outvar = invar.to(dtype=model_dtype), outvar.to(dtype=model_dtype)
+            outvar_pred = model(invar)
+            loss = criterion(outvar_pred, outvar)
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip_norm)
+            torch.cuda.nvtx.range_pop()
+            optimizer.step()
+            scheduler.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}')
+
+            if (j + 1) % cfg.val_freq == 0:
+                model.eval()
+                valid_loss = 0.0
+                with torch.no_grad():
+                    start_time = time.time()
+                    for k, data in enumerate(val_dataloader):
+                        invar = data[0].to(device=local_rank)
+                        outvar = data[1].to(device=local_rank)
+                        cos_zenith = data[2].to(device=local_rank)
+                        in_idx = data[3].item()
+
+                        cos_zenith = torch.squeeze(cos_zenith, dim=2)
+                        cos_zenith = torch.clamp(cos_zenith, min=0.0) - 1.0 / torch.pi  # [b, 2, h, w]
+                        outvar = outvar.to(dtype=model_dtype)
+                        loss = 0.0
+                        for t in range(outvar.shape[1]):
+                            day_of_year, time_of_day = divmod(in_idx + t * cfg.dt, 24 // cfg.dt)
+                            normalized_day_of_year = torch.tensor((day_of_year / 365) * (np.pi / 2), dtype=torch.float32, device=local_rank)
+                            normalized_time_of_day = torch.tensor((time_of_day / (24 - cfg.dt)) * (np.pi / 2), dtype=torch.float32, device=local_rank)
+                            sin_day_of_year = torch.sin(normalized_day_of_year).expand(1, 1, 721, 1440)
+                            cos_day_of_year = torch.cos(normalized_day_of_year).expand(1, 1, 721, 1440)
+                            sin_time_of_day = torch.sin(normalized_time_of_day).expand(1, 1, 721, 1440)
+                            cos_time_of_day = torch.cos(normalized_time_of_day).expand(1, 1, 721, 1440)
+                            invar = torch.concat((invar, cos_zenith, static_data, sin_day_of_year, cos_day_of_year, sin_time_of_day, cos_time_of_day), dim=1)
+                            invar = invar.to(dtype=model_dtype)
+                            outpred = model(invar)
+                            invar = outpred
+                            loss += criterion(outpred, outvar[:, t])
+                        
+                        loss /= outvar.shape[1]
+                        if cfg.world_size > 1:
+                            loss_tensor = loss.detach().to(local_rank)  # torch.tensor(loss, device=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}-{k+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)/(k+1): .02f}s/{cfg_data.dataloader.batch_size}batch] '
+                                    f'loss:{valid_loss / (k+1): .04f}')
+                        
+                    valid_loss /= len(val_dataloader)
+                    is_save_ckp = False
+                    if valid_loss < best_valid_loss:
+                        best_valid_loss = valid_loss
+                        best_loss_epoch = epoch
+                        world_rank == 0 and save_checkpoint(model, optimizer, scheduler, best_valid_loss, best_loss_epoch, cfg.checkpoint_dir)
+                        is_save_ckp = True
+                    
+                train_loss /= (j+1)
+                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)
+                    np.save(train_loss_file, train_losses)
+                
+
+        if epoch - best_loss_epoch > cfg.patience:
+            print(f"Loss has not decrease in {cfg.patience} epochs, stopping training...")
+            exit()
+
+
+def save_checkpoint(model, optimizer, scheduler, best_valid_loss, best_loss_epoch, model_path):
+    model_to_save = model.module if hasattr(model, "module") else model
+    state = {"model_state_dict": model_to_save.state_dict(),
+             "optimizer_state_dict": optimizer.state_dict(),
+             "scheduler_state_dict": scheduler.state_dict(),
+             "best_valid_loss": best_valid_loss,
+             "best_loss_epoch": best_loss_epoch,
+            }
+    torch.save(state, f"{model_path}/model.pth")
+    ### the weight file saving may interrupted due to DCU queue limit, get a backup to ensure there at least has one model 
+    os.system(f"mv {model_path}/model.pth {model_path}/model_bak.pth")
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    main()
+