Upload folder using huggingface_hub

.ms_upload_cache

@@ -0,0 +1 @@
+{"version": 3, "repo_id": "OneScience/Deep-Generative-Model-of-Rain", "files": {"README.md|1787118190.0|6414": {"hash": "ae8b571106cc431550791a52cd3d4c6710031e84363fda4e4a07a3da2ce590d5", "size": 6414, "status": "c"}, "model/__pycache__/dgmr.cpython-311.pyc|1787129401.0|5498": {"hash": "319d638e43312b0c32a8a81e3283ebf22a0c12842183c09091a40a588d08ba69", "size": 5498, "status": "c"}, "conf/config.yaml|1787189067.0|2278": {"hash": "5f07acfc8e3e487a48ebea90399a0207f3ac805a64d16ae7e64da0c749fc2ec8", "size": 2278, "status": "c"}, "model/dgmr_official/__pycache__/__init__.cpython-311.pyc|1787129401.0|506": {"hash": "55bd930076be8ece868e147ca91dbeebee990f54f13d2e470e68221ce5fb631a", "size": 506, "status": "c"}, "LICENSE|1787047051.0|1073": {"hash": "6a9b90043bd7e6836db9a58125c81141a4ee4b8a538aabe33f6dc5180ecff09f", "size": 1073, "status": "c"}, "configuration.json|1787059848.0|39": {"hash": "6b6f30f1e82ad5243c2ff583dcd33b77949442fbf942d80d9f236ffd7a9fe796", "size": 39, "status": "c"}, "model/dgmr.py|1787047183.0|4699": {"hash": "7568f5283b46786a9c034a4cbe5dae727f465bc6184091e239e8f5227bba12a7", "size": 4699, "status": "c"}, "model/dgmr_official/__init__.py|1787047065.0|275": {"hash": "203ab3311b15f33e6130ef33d2da92b820aee85278ba2883ddecdc2af957e790", "size": 275, "status": "c"}, "model/dgmr_official/__pycache__/common.cpython-311.pyc|1787129401.0|20655": {"hash": "52fa87eb304c8ee51b27c65095747576f5b7bf101362b22609f7f4ae87ad299b", "size": 20655, "status": "c"}, "model/dgmr_official/__pycache__/discriminators.cpython-311.pyc|1787129401.0|10887": {"hash": "944d0b3e3fb10fe60264dc396a777d0d10bc3b7a0f55a6676985cb3bfd326d7d", "size": 10887, "status": "c"}, "model/dgmr_official/__pycache__/generators.cpython-311.pyc|1787129402.0|11778": {"hash": "2d7c85d6ed1eab70efce07cc40f227282fd85fa8cd115b79187f59a4a7e7cb81", "size": 11778, "status": "c"}, "model/dgmr_official/__pycache__/losses.cpython-311.pyc|1787129402.0|5091": {"hash": "0e3704bb86d91c76e3a42989c87f622d6c1f989665ea13fe962e00ff10d1efa8", "size": 5091, "status": "c"}, "model/dgmr_official/common.py|1787047241.0|17391": {"hash": "91883f5d9de44bf2821bfa01991fa1b2280f74e2f63089ea1d3c4ad75aea1afa", "size": 17391, "status": "c"}, "model/dgmr_official/discriminators.py|1787047246.0|8891": {"hash": "8f6d93f6008527e3882467d2d57f4244f050076c7bcc633b8dcfb4f33410f260", "size": 8891, "status": "c"}, "model/dgmr_official/generators.py|1787047246.0|7749": {"hash": "1e9ceb835b6d0b7f5aae389169569c29dc8419ccb21c1145f50a8f7c22470102", "size": 7749, "status": "c"}, "model/dgmr_official/layers/Attention.py|1787047051.0|2824": {"hash": "a0e04195b17d9fba45f1b58e37a58524c855f949b86bee52eb370511527e3b57", "size": 2824, "status": "c"}, "model/dgmr_official/layers/ConvGRU.py|1787047051.0|3474": {"hash": "8e97e7df8deeaa3e497365fbb5932b1a7ae2d13cf82656f82e2c40d9985c7e41", "size": 3474, "status": "c"}, "model/dgmr_official/layers/CoordConv.py|1787047051.0|2482": {"hash": "e076d5d5034bcf1bc8a16693fbbf8e348e83b65399c66250dad6b9c6bbcd9915", "size": 2482, "status": "c"}, "model/dgmr_official/layers/__init__.py|1787047051.0|134": {"hash": "37a61b078b8814e0a8c5dce59c246f4cbf83acc6a885a60e7a5afc7bf347d27b", "size": 134, "status": "c"}, "model/dgmr_official/layers/__pycache__/Attention.cpython-311.pyc|1787129401.0|3816": {"hash": "80461302aaefcdbd65d0c30e5bba4e6fbec53f20d2f0b287d3011ba5065b2ac7", "size": 3816, "status": "c"}, "model/dgmr_official/layers/__pycache__/ConvGRU.cpython-311.pyc|1787129401.0|4885": {"hash": "e12ad2510292b13e7a135eeeb59a50f7e4d4221d984fa5c1eccccdb495910f32", "size": 4885, "status": "c"}, "model/dgmr_official/layers/__pycache__/CoordConv.cpython-311.pyc|1787129401.0|4614": {"hash": "870a9883bfc11c72463b31ae5b8c3362147f74c8760870169fbb66067f66f446", "size": 4614, "status": "c"}, "model/dgmr_official/layers/__pycache__/__init__.cpython-311.pyc|1787129401.0|423": {"hash": "ca06f7dc005aca0179c32edcdfbde241b54ca923694ab232a4c1f220b5be73b2", "size": 423, "status": "c"}, "model/dgmr_official/layers/__pycache__/utils.cpython-311.pyc|1787129401.0|1008": {"hash": "dcb03036c637d7fed95d1d5f794857a1d4662a77b1b366a978fc985cb9b8d84e", "size": 1008, "status": "c"}, "model/dgmr_official/layers/utils.py|1787047246.0|528": {"hash": "37926142863e60911d80c7208ee72380a1dccbc759c14101208c6c4abfe3271f", "size": 528, "status": "c"}, "model/dgmr_official/losses.py|1787047133.0|2873": {"hash": "38a72b2b11ca5cd6f2920ac5fbb32fb70bb6c8c6eda7e34a04e381d1ec5571b5", "size": 2873, "status": "c"}, "scripts/fake_data.py|1787118238.0|3066": {"hash": "88dc76ffa492214503a7446d8880234fe14ec88504b1a97bcc28b232f005f8f5", "size": 3066, "status": "c"}, "scripts/inference.py|1787060096.0|3338": {"hash": "e52039c1c3fe77bd130b2379b5d15e8db518e41bb92c206ca8f6f107ea77ecd3", "size": 3338, "status": "c"}, "scripts/train.py|1787060016.0|11345": {"hash": "62c589a0058374338138d065856a6290a839c614cf62133efa4f338621fa761e", "size": 11345, "status": "c"}, "weight/.gitkeep|1787059908.0|1": {"hash": "01ba4719c80b6fe911b091a7c05124b64eeece964e09c058ef8f9805daca546b", "size": 1, "status": "c"}, "scripts/result.py|1787129360.0|8669": {"hash": "2041021d7575eb66e89f2577174a338e6cd2f063826f556a8880bd344306a05c", "size": 8669, "status": "c"}}}

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Open Climate Fix
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,179 @@
+---
+license: mit
+language:
+- en
+- zh
+tags:
+- OneScience
+- Earth science
+- Weather forecasting
+- Short-term precipitation nowcasting
+- Deep generative models
+- Generative adversarial networks
+frameworks: PyTorch
+datasets:
+  - OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">DGMR</span>
+  </strong>
+</p>
+
+
+# Model Introduction
+
+DGMR (Deep Generative Model of Radar) was proposed by DeepMind in 2021. It is a short-term precipitation nowcasting model based on a conditional generative adversarial network (cGAN). The generator consists of a latent conditioning stack, a context conditioning stack, and an autoregressive sampler based on multiple ConvGRU layers. The discriminator evaluates both spatial and temporal dimensions, and the training objective combines hinge GAN loss with a grid-cell regularizer.
+
+Paper:Skillful Precipitation Nowcasting using Deep Generative Models of Radar
+
+https://arxiv.org/abs/2104.00954
+
+# Model Description
+
+DGMR is a probabilistic short-term precipitation nowcasting model: it takes four consecutive radar echo frames and generates 18 future frames at once (5-minute steps, 90 minutes in total, in the paper). Its output is a sample rather than a deterministic estimate. This repository is organized from the official Open Climate Fix PyTorch implementation (`openclimatefix/dgmr`, MIT License) and integrated with the OneScience data loading and training workflow.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Short-term precipitation nowcasting research | Train a cGAN on radar echo sequences to generate future radar echo frames. |
+| Probabilistic forecast output | Obtain multiple generated samples of future fields through latent-variable sampling. |
+| Local quick validation | Use synthetic data to check data loading, training, inference, and result scripts. |
+| ModelScope/OneCode execution | Download the model package, install dependencies, and run the scripts directly. |
+| Multi-card training | Launch multi-process training with `torchrun`. |
+
+
+# Usage
+
+## 1. OneCode Usage
+
+Use the OneCode online environment for intelligent one-click AI4S programming:
+
+[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- GPU or DCU is recommended.
+- CPU can be used for imports and small-configuration connectivity validation, but full training and inference are slower.
+- DCU users must install DTK beforehand. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
+
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/DGMR --local-dir ./DGMR
+cd DGMR
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# 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
+# 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
+
+The OneScience community provides ERA5 data for training (the current repository contains complete data slices subject to data-file size limits). Download it with the command below and confirm that the data path in `conf/config.yaml` is correct:
+
+```bash
+hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data
+```
+
+For a quick workflow validation, run the synthetic data script first:
+
+```bash
+python scripts/fake_data.py
+```
+
+> Note: `scripts/fake_data.py` generates a single-channel radar sequence from `num_context`, `forecast_steps`, batch size, and radar grid size. The current small configuration uses 4 input frames, 6 output frames, and a 128×128 grid.
+
+### Training
+
+Single card：
+
+```bash
+python scripts/train.py
+```
+
+Multiple cards：
+
+```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
+data/checkpoints/valoss.npy
+```
+
+### Training Weights
+The `weight/` folder is reserved for model weights. Pretrained weights are not provided by default; users may train the model using the paper configuration. Official DeepMind pretrained weights require access approval and differ from the current configuration in channel count and grid size, so they must be aligned before loading.
+
+### Inference
+
+Inference reads `data/checkpoints/model_bak.pth`：
+
+```bash
+python scripts/inference.py
+```
+
+Prediction results are written frame by frame to:
+
+```text
+result/output/
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Outputs include：
+
+- `result/rmse.npy`
+- `result/acc.npy`
+- `result/loss.png`
+- Forecast comparison plots for the specified date and variables
+
+
+# Official Source and Reproduction Notes
+
+- The model implementation comes from the official Open Climate Fix PyTorch package `openclimatefix/dgmr` (MIT License). The related network modules (common/layers/generators/discriminators/losses) are embedded unchanged in `model/dgmr_official/` (with the HuggingFace hub mixin and pytorch_lightning training loop removed); `model/dgmr.py` is only a YAML-driven thin wrapper.
+- `conf/config.yaml` uses a small configuration (`forecast_steps=6`, `output_shape=128`, `latent_channels=384`, `context_channels=192`) for connectivity validation by default; paper-level reproduction requires 4→18 frames, a 256×256 grid, `latent_channels=768`, and `context_channels=384`.
+- The paper-level configuration (`num_context=4`, `forecast_steps=18`) requires at least `T >= num_context + forecast_steps + 1 = 23` frames per year; the current synthetic data uses `T=10` (`time_step=6h`), so the default configuration uses 4→6 frames.
+- The discriminator contains BatchNorm1d, so the training batch must be `>= 2` (with synthetic data, the training set uses two samples from two years, exactly one batch).
+- The following details are not disclosed in the paper and are assumptions in this reproduction:data normalization statistics (identity normalization is currently used; real statistics will be supplied with the data), and the number of randomly sampled time frames in the discriminator.
+
+# Official OneScience 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 and License
+
+- This repository is an independent organization and adaptation of DGMR. The model source is based on the Open Climate Fix `openclimatefix/dgmr` implementation and follows the MIT License.
+- Please cite:Ravuri et al. Skilful Precipitation Nowcasting using Deep Generative Models of Radar. Nature 597, 672-677, 2021.

conf/config.yaml

@@ -0,0 +1,65 @@
+# DGMR 训练配置示例
+# 论文配置：4 帧雷达观测上下文（num_context=4）预测 18 帧（forecast_steps=18），
+# 空间 256×256 单通道，latent_channels=768 / context_channels=384（生成器约 13.4M 参数）。
+# 当前默认配置为连通性验证小配置：空间 128×128，预测 6 帧；fake_data.py
+# 根据 num_context、forecast_steps 与 batch_size 自动推导每年所需时间长度 T。
+model:
+  start_epoch: 0
+  max_epoch: 100
+  lr: 1E-4                 # 论文使用 Adam lr=1e-4 分别训练生成器与判别器
+  lr_disc: 1E-4
+  patience: 50
+  checkpoint_dir: "./data/checkpoints"
+
+  # DGMR 结构参数（论文值见注释）
+  forecast_steps: 6        # 预测帧数，论文为 18；小配置取 6
+  num_context: 4           # 上下文（观测）帧数，论文为 4
+  input_channels: 1        # 每帧通道数，论文为 1（雷达回波）
+  output_shape: 128        # 空间尺寸（需 32 整除），论文为 256
+  conv_type: "standard"    # standard / coord / 3d
+  latent_channels: 384     # 论文为 768
+  context_channels: 192    # 论文为 384
+  generation_steps: 6      # 网格单元正则器 Monte-Carlo 采样次数，论文为 6
+  grid_lambda: 20.0        # 网格单元正则器权重，论文为 20
+  precip_weight_cap: 24.0  # 网格单元权重函数上限，论文为 24
+
+# 整个数据读取流程
+datapipe:
+  name: "ERA5"
+  task: "weather_forecasting"
+
+  dataset:
+    type: "hdf5"
+    data_dir: './data/'
+    train_time: [1951, 1952]
+    val_time: [1953]
+    test_time: [1954]
+    img_size: [128, 128]
+    verbose: true
+    cache: false
+
+    # 气象变量（DGMR 面向单通道雷达回波场，此处为占位变量名）
+    channels: ['rain_radar']
+
+  # DataLoader 配置
+  # 注意：DGMR 判别器内部含 BatchNorm1d，训练 batch 必须 >= 2。
+  # fake_data.py 会保证每个年份至少产生 batch_size 个完整时间窗口。
+  dataloader:
+    mask_dtype: "float32"
+    batch_size: 2
+    num_workers: 1
+    pin_memory: true
+    drop_last: true
+    shuffle: false
+    prefetch_factor: 2
+    persistent_workers: true
+
+  # 分布式配置
+  distributed:
+    enabled: true
+    sampler: "DistributedSampler"
+    rank: 0
+    world_size: 2
+    shuffle: true
+    seed: 42
+    drop_last: true

config.json

@@ -0,0 +1,74 @@
+{
+  "model_name": "DGMR",
+  "model_type": "dgmr",
+  "architectures": [
+    "DGMR"
+  ],
+  "framework": "PyTorch",
+  "domain": "climate-and-atmosphere",
+  "task": "precipitation-nowcasting",
+  "implementation": {
+    "entry_point": "model/dgmr.py",
+    "scope": "YAML-configured DGMR conditional GAN wrapper with vendored Open Climate Fix generator, spatial and temporal discriminators, and training losses"
+  },
+  "architecture": {
+    "family": "conditional generative adversarial network for probabilistic radar nowcasting",
+    "generator": "context conditioning stack plus latent conditioning stack and a four-scale ConvGRU autoregressive sampler",
+    "discriminator": "combined spatial and temporal discriminators with spectrally normalized residual blocks",
+    "training_objective": "hinge GAN loss plus weighted grid-cell regularizer",
+    "input_format": "BTCHW",
+    "output_format": "BTCHW",
+    "activation": "ReLU",
+    "normalization": "BatchNorm",
+    "repository_default_config": {
+      "purpose": "small connectivity-validation configuration",
+      "num_context": 4,
+      "forecast_steps": 6,
+      "input_channels": 1,
+      "output_shape": 128,
+      "conv_type": "standard",
+      "latent_channels": 384,
+      "context_channels": 192,
+      "generation_steps": 6,
+      "grid_lambda": 20.0,
+      "precip_weight_cap": 24.0
+    },
+    "paper_configuration": {
+      "num_context": 4,
+      "forecast_steps": 18,
+      "input_channels": 1,
+      "output_shape": 256,
+      "latent_channels": 768,
+      "context_channels": 384,
+      "generation_steps": 6,
+      "grid_lambda": 20.0,
+      "precip_weight_cap": 24.0
+    }
+  },
+  "data": {
+    "dataset": "ERA5Datapipe-compatible HDF5 placeholder radar sequences",
+    "variable": "rain_radar",
+    "input_length": 4,
+    "output_length": 6,
+    "channels": 1,
+    "default_smoke_spatial_size": [
+      128,
+      128
+    ],
+    "paper_frame_interval_minutes": 5,
+    "paper_output_length": 18,
+    "paper_spatial_size": [
+      256,
+      256
+    ],
+    "normalization": "identity statistics in synthetic data; real-data statistics are read from HDF5"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/dgmr.py",
+    "model/dgmr_official/",
+    "scripts/train.py",
+    "scripts/fake_data.py",
+    "README.md"
+  ]
+}

configuration.json

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

model/dgmr.py

@@ -0,0 +1,114 @@
+# coding=utf-8
+#
+# SPDX-License-Identifier: MIT
+#
+# Copyright (c) 2021 Open Climate Fix
+#
+# This module is a thin configuration wrapper around the DGMR (Deep
+# Generative Model of Radar) architecture from Ravuri et al. (2021,
+# "Skilful Precipitation Nowcasting using Deep Generative Models of Radar",
+# Nature 597), as re-implemented in PyTorch by Open Climate Fix
+# (``openclimatefix/skillful_nowcasting``, MIT License). The network modules
+# are vendored verbatim (minus HuggingFace hub mixins) under
+# ``dgmr_official/``; only config plumbing and the plain ``forward`` are
+# added here for YAML-driven usage.
+import torch
+import torch.nn as nn
+
+from model.dgmr_official.common import ContextConditioningStack, LatentConditioningStack
+from model.dgmr_official.discriminators import Discriminator
+from model.dgmr_official.generators import Generator, Sampler
+from model.dgmr_official.losses import GridCellLoss
+
+
+def weight_fn(y, precip_weight_cap=24.0):
+    """
+    Weight function for the grid cell loss: w(y) = max(y + 1, cap).
+    """
+    return torch.max(y + 1, torch.tensor(precip_weight_cap, device=y.device))
+
+
+class DGMR(nn.Module):
+    """
+    Config-driven DGMR wrapper (generator + discriminator).
+
+    The generator is a conditional GAN generator that takes ``num_context``
+    observed radar frames of shape [B, T, C, H, W] and produces
+    ``forecast_steps`` future frames of the same spatial size. The
+    discriminator scores full sequences (context + forecast) spatially and
+    temporally; during GAN training the hinge losses plus the grid-cell
+    regularizer are applied (see ``dgmr_official/losses.py`` and the paper).
+
+    Args:
+        forecast_steps: Number of frames to predict in the future (paper: 18).
+        num_context: Number of input/context frames (paper: 4).
+        input_channels: Number of channels per frame (paper: 1, radar).
+        output_shape: Spatial size of the frames; must be divisible by 32
+            (paper: 256). Discriminators additionally need >= 128 px.
+        conv_type: Convolution flavour used by the conditioning stack,
+            one of "standard" / "coord" / "3d".
+        latent_channels / context_channels: DGMR architecture sizes
+            (paper: 768 / 384).
+        generation_steps: Number of Monte-Carlo generator samples used when
+            computing the grid-cell regularizer during training (paper: 6).
+        grid_lambda: Weight of the grid-cell regularizer (paper: 20).
+        precip_weight_cap: Ceiling for the grid-cell weight function (paper: 24).
+    """
+
+    def __init__(
+        self,
+        forecast_steps: int = 18,
+        num_context: int = 4,
+        input_channels: int = 1,
+        output_shape: int = 256,
+        conv_type: str = "standard",
+        latent_channels: int = 768,
+        context_channels: int = 384,
+        generation_steps: int = 6,
+        grid_lambda: float = 20.0,
+        precip_weight_cap: float = 24.0,
+    ):
+        super().__init__()
+        self.forecast_steps = int(forecast_steps)
+        self.num_context = int(num_context)
+        self.input_channels = int(input_channels)
+        self.output_shape = int(output_shape)
+        self.conv_type = conv_type
+        self.latent_channels = int(latent_channels)
+        self.context_channels = int(context_channels)
+        self.generation_steps = int(generation_steps)
+        self.grid_lambda = float(grid_lambda)
+        self.precip_weight_cap = float(precip_weight_cap)
+
+        self.conditioning_stack = ContextConditioningStack(
+            input_channels=self.input_channels,
+            conv_type=self.conv_type,
+            output_channels=self.context_channels,
+        )
+        self.latent_stack = LatentConditioningStack(
+            shape=(
+                8 * self.input_channels,
+                self.output_shape // 32,
+                self.output_shape // 32,
+            ),
+            output_channels=self.latent_channels,
+        )
+        self.sampler = Sampler(
+            forecast_steps=self.forecast_steps,
+            latent_channels=self.latent_channels,
+            context_channels=self.context_channels,
+        )
+        self.generator = Generator(self.conditioning_stack, self.latent_stack, self.sampler)
+        self.discriminator = Discriminator(self.input_channels)
+        self.grid_regularizer = GridCellLoss(
+            weight_fn=weight_fn, precip_weight_cap=self.precip_weight_cap
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Observed radar frames, shape [batch, num_context, C, H, W].
+        Returns:
+            Forecast frames, shape [batch, forecast_steps, C, H, W].
+        """
+        return self.generator(x)

model/dgmr_official/__init__.py

@@ -0,0 +1,6 @@
+"""Vendored DGMR (Deep Generative Model of Radar) modules.
+
+Stripped from the ``dgmr`` PyTorch package (Open Climate Fix, MIT License):
+the ``PyTorchModelHubMixin`` HuggingFace mixin and the Lightning training
+loop are removed; the model modules themselves are verbatim.
+"""

model/dgmr_official/common.py

@@ -0,0 +1,497 @@
+"""Modules for generator blocks."""
+
+from typing import Tuple
+
+import einops
+import torch
+import torch.nn.functional as F
+
+from torch.distributions import normal
+from torch.nn.modules.pixelshuffle import PixelUnshuffle
+from torch.nn.utils.parametrizations import spectral_norm
+
+from .layers import AttentionLayer
+from .layers.utils import get_conv_layer
+
+
+class GBlock(torch.nn.Module):
+    """Residual generator block without upsampling."""
+
+    def __init__(
+        self,
+        input_channels: int = 12,
+        output_channels: int = 12,
+        conv_type: str = "standard",
+        spectral_normalized_eps=0.0001,
+    ):
+        """
+        G Block from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf.
+
+        Args:
+            input_channels: Number of input channels
+            output_channels: Number of output channels
+            conv_type: Type of convolution desired, see satflow/models/utils.py for options
+            spectral_normalized_eps: constrains the spectral norm of the weights.
+        """
+        super().__init__()
+        self.output_channels = output_channels
+        self.bn1 = torch.nn.BatchNorm2d(input_channels)
+        self.bn2 = torch.nn.BatchNorm2d(input_channels)
+        self.relu = torch.nn.ReLU()
+        # Upsample in the 1x1
+        conv2d = get_conv_layer(conv_type)
+        self.conv_1x1 = spectral_norm(
+            conv2d(
+                in_channels=input_channels,
+                out_channels=output_channels,
+                kernel_size=1,
+            ),
+            eps=spectral_normalized_eps,
+        )
+        # Upsample 2D conv
+        self.first_conv_3x3 = spectral_norm(
+            conv2d(
+                in_channels=input_channels,
+                out_channels=input_channels,
+                kernel_size=3,
+                padding=1,
+            ),
+            eps=spectral_normalized_eps,
+        )
+        self.last_conv_3x3 = spectral_norm(
+            conv2d(
+                in_channels=input_channels, out_channels=output_channels, kernel_size=3, padding=1
+            ),
+            eps=spectral_normalized_eps,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the forward function."""
+        # Optionally spectrally normalized 1x1 convolution
+        if x.shape[1] != self.output_channels:
+            sc = self.conv_1x1(x)
+        else:
+            sc = x
+
+        x2 = self.bn1(x)
+        x2 = self.relu(x2)
+        x2 = self.first_conv_3x3(x2)  # Make sure size is doubled
+        x2 = self.bn2(x2)
+        x2 = self.relu(x2)
+        x2 = self.last_conv_3x3(x2)
+        # Sum combine, residual connection
+        x = x2 + sc
+        return x
+
+
+class UpsampleGBlock(torch.nn.Module):
+    """Residual generator block with upsampling."""
+
+    def __init__(
+        self,
+        input_channels: int = 12,
+        output_channels: int = 12,
+        conv_type: str = "standard",
+        spectral_normalized_eps=0.0001,
+    ):
+        """
+        G Block from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf.
+
+        Args:
+            input_channels: Number of input channels.
+            output_channels: Number of output channels.
+            conv_type: Type of convolution desired, see satflow/models/utils.py for options.
+            spectral_normalized_eps: constrains the spectral norm of the weights.
+        """
+        super().__init__()
+        self.output_channels = output_channels
+        self.bn1 = torch.nn.BatchNorm2d(input_channels)
+        self.bn2 = torch.nn.BatchNorm2d(input_channels)
+        self.relu = torch.nn.ReLU()
+        # Upsample in the 1x1
+        conv2d = get_conv_layer(conv_type)
+        self.conv_1x1 = spectral_norm(
+            conv2d(
+                in_channels=input_channels,
+                out_channels=output_channels,
+                kernel_size=1,
+            ),
+            eps=spectral_normalized_eps,
+        )
+        self.upsample = torch.nn.Upsample(scale_factor=2, mode="nearest")
+        # Upsample 2D conv
+        self.first_conv_3x3 = spectral_norm(
+            conv2d(
+                in_channels=input_channels,
+                out_channels=input_channels,
+                kernel_size=3,
+                padding=1,
+            ),
+            eps=spectral_normalized_eps,
+        )
+        self.last_conv_3x3 = spectral_norm(
+            conv2d(
+                in_channels=input_channels, out_channels=output_channels, kernel_size=3, padding=1
+            ),
+            eps=spectral_normalized_eps,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the forward function."""
+        # Spectrally nsormalized 1x1 convolution
+        sc = self.upsample(x)
+        sc = self.conv_1x1(sc)
+
+        x2 = self.bn1(x)
+        x2 = self.relu(x2)
+        # Upsample
+        x2 = self.upsample(x2)
+        x2 = self.first_conv_3x3(x2)  # Make sure size is doubled
+        x2 = self.bn2(x2)
+        x2 = self.relu(x2)
+        x2 = self.last_conv_3x3(x2)
+        # Sum combine, residual connection
+        x = x2 + sc
+        return x
+
+
+class DBlock(torch.nn.Module):
+    """D block class."""
+
+    def __init__(
+        self,
+        input_channels: int = 12,
+        output_channels: int = 12,
+        conv_type: str = "standard",
+        first_relu: bool = True,
+        keep_same_output: bool = False,
+    ):
+        """
+        D and 3D Block from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf.
+
+        Args:
+            input_channels: Number of input channels
+            output_channels: Number of output channels
+            conv_type: Convolution type, see satflow/models/utils.py for options
+            first_relu: Whether to have an ReLU before the first 3x3 convolution
+            keep_same_output: Whether the output should have the same spatial dimensions
+            as input, if False, downscales by 2
+        """
+        super().__init__()
+        self.input_channels = input_channels
+        self.output_channels = output_channels
+        self.first_relu = first_relu
+        self.keep_same_output = keep_same_output
+        self.conv_type = conv_type
+        conv2d = get_conv_layer(conv_type)
+        if conv_type == "3d":
+            # 3D Average pooling
+            self.pooling = torch.nn.AvgPool3d(kernel_size=2, stride=2)
+        else:
+            self.pooling = torch.nn.AvgPool2d(kernel_size=2, stride=2)
+        self.conv_1x1 = spectral_norm(
+            conv2d(
+                in_channels=input_channels,
+                out_channels=output_channels,
+                kernel_size=1,
+            )
+        )
+        self.first_conv_3x3 = spectral_norm(
+            conv2d(
+                in_channels=input_channels,
+                out_channels=output_channels,
+                kernel_size=3,
+                padding=1,
+            )
+        )
+        self.last_conv_3x3 = spectral_norm(
+            conv2d(
+                in_channels=output_channels,
+                out_channels=output_channels,
+                kernel_size=3,
+                padding=1,
+                stride=1,
+            )
+        )
+        # Downsample at end of 3x3
+        self.relu = torch.nn.ReLU()
+        # Concatenate to double final channels and keep reduced spatial extent
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the D residual block."""
+        if self.input_channels != self.output_channels:
+            x1 = self.conv_1x1(x)
+            if not self.keep_same_output:
+                x1 = self.pooling(x1)
+        else:
+            x1 = x
+
+        if self.first_relu:
+            x = self.relu(x)
+        x = self.first_conv_3x3(x)
+        x = self.relu(x)
+        x = self.last_conv_3x3(x)
+
+        if not self.keep_same_output:
+            x = self.pooling(x)
+        x = x1 + x  # Sum the outputs should be half spatial and double channels
+        return x
+
+
+class LBlock(torch.nn.Module):
+    """Residual block for the Latent Stack."""
+
+    def __init__(
+        self,
+        input_channels: int = 12,
+        output_channels: int = 12,
+        kernel_size: int = 3,
+        conv_type: str = "standard",
+    ):
+        """
+        Initialize the L-block.
+
+        L-Block for increasing the number of channels in the input
+         from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf
+        Args:
+            input_channels: Number of input channels
+            output_channels: Number of output channels
+            conv_type: Which type of convolution desired, see satflow/models/utils.py for options
+        """
+        super().__init__()
+        # Output size should be channel_out - channel_in
+        self.input_channels = input_channels
+        self.output_channels = output_channels
+        conv2d = get_conv_layer(conv_type)
+        self.conv_1x1 = conv2d(
+            in_channels=input_channels,
+            out_channels=output_channels - input_channels,
+            kernel_size=1,
+        )
+
+        self.first_conv_3x3 = conv2d(
+            input_channels,
+            out_channels=output_channels,
+            kernel_size=kernel_size,
+            padding=1,
+            stride=1,
+        )
+        self.relu = torch.nn.ReLU()
+        self.last_conv_3x3 = conv2d(
+            in_channels=output_channels,
+            out_channels=output_channels,
+            kernel_size=kernel_size,
+            padding=1,
+            stride=1,
+        )
+
+    def forward(self, x) -> torch.Tensor:
+        """Apply the L residual block to this tensor."""
+        if self.input_channels < self.output_channels:
+            sc = self.conv_1x1(x)
+            sc = torch.cat([x, sc], dim=1)
+        else:
+            sc = x
+
+        x2 = self.relu(x)
+        x2 = self.first_conv_3x3(x2)
+        x2 = self.relu(x2)
+        x2 = self.last_conv_3x3(x2)
+        return x2 + sc
+
+
+class ContextConditioningStack(torch.nn.Module):
+    """Context conditioning stack."""
+
+    def __init__(
+        self,
+        input_channels: int = 1,
+        output_channels: int = 768,
+        num_context_steps: int = 4,
+        conv_type: str = "standard",
+    ):
+        """
+        Conditioning Stack using the context images from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf.
+
+        Args:
+            input_channels: Number of input channels per timestep
+            output_channels: Number of output channels for the lowest block
+            num_context_steps: number of context steps (int)
+            conv_type: Type of 2D convolution to use, see satflow/models/utils.py for options
+            **kwargs: Allow initialize of the parameters above through key pairs
+        """
+        super().__init__()
+
+        conv2d = get_conv_layer(conv_type)
+        self.space2depth = PixelUnshuffle(downscale_factor=2)
+        # Process each observation processed separately with 4 downsample blocks
+        # Concatenate across channel dimension, and for each output, 3x3 spectrally
+        # normalized convolution to reduce number of channels by 2, followed by ReLU
+        self.d1 = DBlock(
+            input_channels=4 * input_channels,
+            output_channels=((output_channels // 4) * input_channels) // num_context_steps,
+            conv_type=conv_type,
+        )
+        self.d2 = DBlock(
+            input_channels=((output_channels // 4) * input_channels) // num_context_steps,
+            output_channels=((output_channels // 2) * input_channels) // num_context_steps,
+            conv_type=conv_type,
+        )
+        self.d3 = DBlock(
+            input_channels=((output_channels // 2) * input_channels) // num_context_steps,
+            output_channels=(output_channels * input_channels) // num_context_steps,
+            conv_type=conv_type,
+        )
+        self.d4 = DBlock(
+            input_channels=(output_channels * input_channels) // num_context_steps,
+            output_channels=(output_channels * 2 * input_channels) // num_context_steps,
+            conv_type=conv_type,
+        )
+        self.conv1 = spectral_norm(
+            conv2d(
+                in_channels=(output_channels // 4) * input_channels,
+                out_channels=(output_channels // 8) * input_channels,
+                kernel_size=3,
+                padding=1,
+            )
+        )
+
+        self.conv2 = spectral_norm(
+            conv2d(
+                in_channels=(output_channels // 2) * input_channels,
+                out_channels=(output_channels // 4) * input_channels,
+                kernel_size=3,
+                padding=1,
+            )
+        )
+
+        self.conv3 = spectral_norm(
+            conv2d(
+                in_channels=output_channels * input_channels,
+                out_channels=(output_channels // 2) * input_channels,
+                kernel_size=3,
+                padding=1,
+            )
+        )
+
+        self.conv4 = spectral_norm(
+            conv2d(
+                in_channels=output_channels * 2 * input_channels,
+                out_channels=output_channels * input_channels,
+                kernel_size=3,
+                padding=1,
+            )
+        )
+
+        self.relu = torch.nn.ReLU()
+
+    def forward(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Generate the condition representation."""
+        # Each timestep processed separately
+        x = self.space2depth(x)
+        steps = x.size(1)  # Number of timesteps
+        scale_1 = []
+        scale_2 = []
+        scale_3 = []
+        scale_4 = []
+        for i in range(steps):
+            s1 = self.d1(x[:, i, :, :, :])
+            s2 = self.d2(s1)
+            s3 = self.d3(s2)
+            s4 = self.d4(s3)
+            scale_1.append(s1)
+            scale_2.append(s2)
+            scale_3.append(s3)
+            scale_4.append(s4)
+        scale_1 = torch.stack(scale_1, dim=1)  # B, T, C, H, W and want along C dimension
+        scale_2 = torch.stack(scale_2, dim=1)  # B, T, C, H, W and want along C dimension
+        scale_3 = torch.stack(scale_3, dim=1)  # B, T, C, H, W and want along C dimension
+        scale_4 = torch.stack(scale_4, dim=1)  # B, T, C, H, W and want along C dimension
+        # Mixing layer
+        scale_1 = self._mixing_layer(scale_1, self.conv1)
+        scale_2 = self._mixing_layer(scale_2, self.conv2)
+        scale_3 = self._mixing_layer(scale_3, self.conv3)
+        scale_4 = self._mixing_layer(scale_4, self.conv4)
+        return scale_1, scale_2, scale_3, scale_4
+
+    def _mixing_layer(self, inputs, conv_block):
+        """Combine the inputs and then passed into the convolution stack."""
+        # Convert from [batch_size, time, h, w, c] -> [batch_size, h, w, c * time]
+        # then perform convolution on the output while preserving number of c.
+        stacked_inputs = einops.rearrange(inputs, "b t c h w -> b (c t) h w")
+        return F.relu(conv_block(stacked_inputs))
+
+
+class LatentConditioningStack(torch.nn.Module):
+    """Latent conditioning stack class."""
+
+    def __init__(
+        self,
+        shape: (int, int, int) = (8, 8, 8),
+        output_channels: int = 768,
+        use_attention: bool = True,
+    ):
+        """
+        Latent conditioning stack from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf.
+
+        Args:
+            shape: Shape of the latent space, Should be (H/32,W/32,x) of the final image shape
+            output_channels: Number of output channels for the conditioning stack
+            use_attention: Whether to have a self-attention block or not
+            **kwargs: allow initialize of the parameters above through key pairs
+        """
+        super().__init__()
+
+        self.shape = shape
+        self.use_attention = use_attention
+        self.distribution = normal.Normal(loc=torch.Tensor([0.0]), scale=torch.Tensor([1.0]))
+
+        self.conv_3x3 = spectral_norm(
+            torch.nn.Conv2d(
+                in_channels=shape[0], out_channels=shape[0], kernel_size=(3, 3), padding=1
+            )
+        )
+        self.l_block1 = LBlock(input_channels=shape[0], output_channels=output_channels // 32)
+        self.l_block2 = LBlock(
+            input_channels=output_channels // 32, output_channels=output_channels // 16
+        )
+        self.l_block3 = LBlock(
+            input_channels=output_channels // 16, output_channels=output_channels // 4
+        )
+        if self.use_attention:
+            self.att_block = AttentionLayer(
+                input_channels=output_channels // 4, output_channels=output_channels // 4
+            )
+        self.l_block4 = LBlock(input_channels=output_channels // 4, output_channels=output_channels)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Apply convolution, l blocks and spatial attention module to the tensor.
+
+        Args:
+            x: tensor on the correct device, to move over the latent distribution
+
+        Returns:
+               tensor
+
+        """
+        # Independent draws from Norma ldistribution
+        z = self.distribution.sample(self.shape)
+        # Batch is at end for some reason, reshape
+        z = torch.permute(z, (3, 0, 1, 2)).type_as(x)
+
+        # 3x3 Convolution
+        z = self.conv_3x3(z)
+
+        # 3 L Blocks to increase number of channels
+        z = self.l_block1(z)
+        z = self.l_block2(z)
+        z = self.l_block3(z)
+        # Spatial attention module
+        z = self.att_block(z)
+
+        # L block to increase number of channel to 768
+        z = self.l_block4(z)
+        return z

model/dgmr_official/discriminators.py

@@ -0,0 +1,232 @@
+"""Discriminators."""
+
+import torch
+import torch.nn.functional as F
+
+from torch.nn.modules.pixelshuffle import PixelUnshuffle
+from torch.nn.utils.parametrizations import spectral_norm
+
+from .common import DBlock
+
+
+class Discriminator(torch.nn.Module):
+    """Discriminators class."""
+
+    def __init__(
+        self,
+        input_channels: int = 12,
+        num_spatial_frames: int = 8,
+        conv_type: str = "standard",
+    ):
+        """
+        Initialize the discriminator.
+
+        Args:
+            input_channels: Number of input channels (int)
+            num_spatial_frames: Number of spatial frames (int)
+            conv_type: the specified convolution type (str)
+            **kwargs: allow initialize of the parameters above through key pairs
+        """
+        super().__init__()
+
+        self.spatial_discriminator = SpatialDiscriminator(
+            input_channels=input_channels, num_timesteps=num_spatial_frames, conv_type=conv_type
+        )
+        self.temporal_discriminator = TemporalDiscriminator(
+            input_channels=input_channels, conv_type=conv_type
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Mixes the spatial loss and temporal loss of the tensor prior to returning it."""
+        spatial_loss = self.spatial_discriminator(x)
+        temporal_loss = self.temporal_discriminator(x)
+
+        return torch.cat([spatial_loss, temporal_loss], dim=1)
+
+
+class TemporalDiscriminator(torch.nn.Module):
+    """Temporal Discriminator class."""
+
+    def __init__(
+        self,
+        input_channels: int = 12,
+        num_layers: int = 3,
+        conv_type: str = "standard",
+    ):
+        """
+        Temporal Discriminator from the Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf.
+
+        Args:
+            input_channels: Number of channels per timestep
+            crop_size: Size of the crop, in the paper half the width of the input images
+            num_layers: Number of intermediate DBlock layers to use
+            conv_type: Type of 2d convolutions to use, see satflow/models/utils.py for options
+            **kwargs: allow initialize of the parameters above through key pairs
+        """
+        super().__init__()
+
+        self.downsample = torch.nn.AvgPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2))
+        self.space2depth = PixelUnshuffle(downscale_factor=2)
+        internal_chn = 48
+        self.d1 = DBlock(
+            input_channels=4 * input_channels,
+            output_channels=internal_chn * input_channels,
+            conv_type="3d",
+            first_relu=False,
+        )
+        self.d2 = DBlock(
+            input_channels=internal_chn * input_channels,
+            output_channels=2 * internal_chn * input_channels,
+            conv_type="3d",
+        )
+        self.intermediate_dblocks = torch.nn.ModuleList()
+        for _ in range(num_layers):
+            internal_chn *= 2
+            self.intermediate_dblocks.append(
+                DBlock(
+                    input_channels=internal_chn * input_channels,
+                    output_channels=2 * internal_chn * input_channels,
+                    conv_type=conv_type,
+                )
+            )
+
+        self.d_last = DBlock(
+            input_channels=2 * internal_chn * input_channels,
+            output_channels=2 * internal_chn * input_channels,
+            keep_same_output=True,
+            conv_type=conv_type,
+        )
+
+        self.fc = spectral_norm(torch.nn.Linear(2 * internal_chn * input_channels, 1))
+        self.relu = torch.nn.ReLU()
+        self.bn = torch.nn.BatchNorm1d(2 * internal_chn * input_channels)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the forward function."""
+        x = self.downsample(x)
+
+        x = self.space2depth(x)
+        # Have to move time and channels
+        x = torch.permute(x, dims=(0, 2, 1, 3, 4))
+        # 2 residual 3D blocks to halve resolution if image, double number of channels and reduce
+        # number of time steps
+        x = self.d1(x)
+        x = self.d2(x)
+        # Convert back to T x C x H x W
+        x = torch.permute(x, dims=(0, 2, 1, 3, 4))
+        # Per Timestep part now, same as spatial discriminator
+        representations = []
+        for idx in range(x.size(1)):
+            # Intermediate DBlocks
+            # Three residual D Blocks to halve the resolution of the image and double
+            # the number of channels.
+            rep = x[:, idx, :, :, :]
+            for d in self.intermediate_dblocks:
+                rep = d(rep)
+            # One more D Block without downsampling or increase number of channels
+            rep = self.d_last(rep)
+
+            rep = torch.sum(F.relu(rep), dim=[2, 3])
+            rep = self.bn(rep)
+            rep = self.fc(rep)
+
+            representations.append(rep)
+        # The representations are summed together before the ReLU
+        x = torch.stack(representations, dim=1)
+        # Should be [Batch, N, 1]
+        x = torch.sum(x, keepdim=True, dim=1)
+        return x
+
+
+class SpatialDiscriminator(torch.nn.Module):
+    """Spatial Discriminator class."""
+
+    def __init__(
+        self,
+        input_channels: int = 12,
+        num_timesteps: int = 8,
+        num_layers: int = 4,
+        conv_type: str = "standard",
+    ):
+        """
+        Spatial discriminator from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf.
+
+        Args:
+            input_channels: Number of input channels per timestep
+            num_timesteps: Number of timesteps to use, in the paper 8/18 timesteps were chosen
+            num_layers: Number of intermediate DBlock layers to use
+            conv_type: Type of 2d convolutions to use, see satflow/models/utils.py for options
+            **kwargs: allow initialize of the parameters above through key pairs
+        """
+        super().__init__()
+        # Randomly, uniformly, select 8 timesteps to do this on from the input
+        self.num_timesteps = num_timesteps
+        # First step is mean pooling 2x2 to reduce input by half
+        self.mean_pool = torch.nn.AvgPool2d(2)
+        self.space2depth = PixelUnshuffle(downscale_factor=2)
+        internal_chn = 24
+        self.d1 = DBlock(
+            input_channels=4 * input_channels,
+            output_channels=2 * internal_chn * input_channels,
+            first_relu=False,
+            conv_type=conv_type,
+        )
+        self.intermediate_dblocks = torch.nn.ModuleList()
+        for _ in range(num_layers):
+            internal_chn *= 2
+            self.intermediate_dblocks.append(
+                DBlock(
+                    input_channels=internal_chn * input_channels,
+                    output_channels=2 * internal_chn * input_channels,
+                    conv_type=conv_type,
+                )
+            )
+        self.d6 = DBlock(
+            input_channels=2 * internal_chn * input_channels,
+            output_channels=2 * internal_chn * input_channels,
+            keep_same_output=True,
+            conv_type=conv_type,
+        )
+
+        # Spectrally normalized linear layer for binary classification
+        self.fc = spectral_norm(torch.nn.Linear(2 * internal_chn * input_channels, 1))
+        self.relu = torch.nn.ReLU()
+        self.bn = torch.nn.BatchNorm1d(2 * internal_chn * input_channels)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the forward function to the tensor."""
+        # x should be the chosen 8 or so
+        idxs = torch.randint(low=0, high=x.size()[1], size=(self.num_timesteps,))
+        representations = []
+        for idx in idxs:
+            rep = self.mean_pool(x[:, idx, :, :, :])  # 128x128
+            rep = self.space2depth(rep)  # 64x64x4
+            rep = self.d1(rep)  # 32x32
+            # Intermediate DBlocks
+            for d in self.intermediate_dblocks:
+                rep = d(rep)
+            rep = self.d6(rep)  # 2x2
+            rep = torch.sum(F.relu(rep), dim=[2, 3])
+            rep = self.bn(rep)
+            rep = self.fc(rep)
+            """
+            Pseudocode from DeepMind
+            # Sum-pool the representations and feed to spectrally normalized lin. layer.
+            y = tf.reduce_sum(tf.nn.relu(y), axis=[1, 2])
+            y = layers.BatchNorm(calc_sigma=False)(y)
+            output_layer = layers.Linear(output_size=1)
+            output = output_layer(y)
+
+            # Take the sum across the t samples. Note: we apply the ReLU to
+            # (1 - score_real) and (1 + score_generated) in the loss.
+            output = tf.reshape(output, [b, n, 1])
+            output = tf.reduce_sum(output, keepdims=True, axis=1)
+            return output
+            """
+            representations.append(rep)
+
+        # The representations are summed together before the ReLU
+        x = torch.stack(representations, dim=1)
+        # Should be [Batch, N, 1]
+        x = torch.sum(x, keepdim=True, dim=1)
+        return x

model/dgmr_official/generators.py

@@ -0,0 +1,212 @@
+"""Generators."""
+
+import logging
+from typing import List
+
+import einops
+import torch
+import torch.nn.functional as F
+
+from torch.nn.modules.pixelshuffle import PixelShuffle
+from torch.nn.utils.parametrizations import spectral_norm
+
+from .common import GBlock, UpsampleGBlock
+from .layers import ConvGRU
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARN)
+
+
+class Sampler(torch.nn.Module):
+    """Sampler class."""
+
+    def __init__(
+        self,
+        forecast_steps: int = 18,
+        latent_channels: int = 768,
+        context_channels: int = 384,
+        output_channels: int = 1,
+    ):
+        """
+        Sampler from the Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf.
+
+        The sampler takes the output from the Latent and Context conditioning stacks and
+        creates one stack of ConvGRU layers per future timestep.
+
+        Args:
+            forecast_steps: Number of forecast steps (int)
+            latent_channels: Number of input channels to the lowest ConvGRU layer (int)
+            context_channels: Number of context channels (int)
+            output_channels: Number of output channels (int)
+            **kwargs: allow initialize of the parameters above through key pairs
+        """
+        super().__init__()
+
+        self.forecast_steps = forecast_steps
+
+        self.convGRU1 = ConvGRU(
+            input_channels=latent_channels + context_channels,
+            output_channels=context_channels,
+            kernel_size=3,
+        )
+        self.gru_conv_1x1 = spectral_norm(
+            torch.nn.Conv2d(
+                in_channels=context_channels, out_channels=latent_channels, kernel_size=(1, 1)
+            )
+        )
+        self.g1 = GBlock(input_channels=latent_channels, output_channels=latent_channels)
+        self.up_g1 = UpsampleGBlock(
+            input_channels=latent_channels, output_channels=latent_channels // 2
+        )
+
+        self.convGRU2 = ConvGRU(
+            input_channels=latent_channels // 2 + context_channels // 2,
+            output_channels=context_channels // 2,
+            kernel_size=3,
+        )
+        self.gru_conv_1x1_2 = spectral_norm(
+            torch.nn.Conv2d(
+                in_channels=context_channels // 2,
+                out_channels=latent_channels // 2,
+                kernel_size=(1, 1),
+            )
+        )
+        self.g2 = GBlock(input_channels=latent_channels // 2, output_channels=latent_channels // 2)
+        self.up_g2 = UpsampleGBlock(
+            input_channels=latent_channels // 2, output_channels=latent_channels // 4
+        )
+
+        self.convGRU3 = ConvGRU(
+            input_channels=latent_channels // 4 + context_channels // 4,
+            output_channels=context_channels // 4,
+            kernel_size=3,
+        )
+        self.gru_conv_1x1_3 = spectral_norm(
+            torch.nn.Conv2d(
+                in_channels=context_channels // 4,
+                out_channels=latent_channels // 4,
+                kernel_size=(1, 1),
+            )
+        )
+        self.g3 = GBlock(input_channels=latent_channels // 4, output_channels=latent_channels // 4)
+        self.up_g3 = UpsampleGBlock(
+            input_channels=latent_channels // 4, output_channels=latent_channels // 8
+        )
+
+        self.convGRU4 = ConvGRU(
+            input_channels=latent_channels // 8 + context_channels // 8,
+            output_channels=context_channels // 8,
+            kernel_size=3,
+        )
+        self.gru_conv_1x1_4 = spectral_norm(
+            torch.nn.Conv2d(
+                in_channels=context_channels // 8,
+                out_channels=latent_channels // 8,
+                kernel_size=(1, 1),
+            )
+        )
+        self.g4 = GBlock(input_channels=latent_channels // 8, output_channels=latent_channels // 8)
+        self.up_g4 = UpsampleGBlock(
+            input_channels=latent_channels // 8, output_channels=latent_channels // 16
+        )
+
+        self.bn = torch.nn.BatchNorm2d(latent_channels // 16)
+        self.relu = torch.nn.ReLU()
+        self.conv_1x1 = spectral_norm(
+            torch.nn.Conv2d(
+                in_channels=latent_channels // 16,
+                out_channels=4 * output_channels,
+                kernel_size=(1, 1),
+            )
+        )
+
+        self.depth2space = PixelShuffle(upscale_factor=2)
+
+    def forward(
+        self, conditioning_states: List[torch.Tensor], latent_dim: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Perform the sampling from Skillful Nowcasting with GANs.
+
+        Args:
+            conditioning_states: Outputs from the `ContextConditioningStack` with the 4 input
+            states, ordered from largest to smallest spatially latent_dim: Output from
+            `LatentConditioningStack` for input into the ConvGRUs
+
+            latent_dim: (torch.Tensor)
+
+        Returns:
+            forecast_steps-length output of images for future timesteps
+
+        """
+        # Iterate through each forecast step
+        # Initialize with conditioning state for first one, output for second one
+        init_states = conditioning_states
+        # Expand latent dim to match batch size
+        latent_dim = einops.repeat(
+            latent_dim, "b c h w -> (repeat b) c h w", repeat=init_states[0].shape[0]
+        )
+        hidden_states = [latent_dim] * self.forecast_steps
+
+        # Layer 4 (bottom most)
+        hidden_states = self.convGRU1(hidden_states, init_states[3])
+        hidden_states = [self.gru_conv_1x1(h) for h in hidden_states]
+        hidden_states = [self.g1(h) for h in hidden_states]
+        hidden_states = [self.up_g1(h) for h in hidden_states]
+
+        # Layer 3.
+        hidden_states = self.convGRU2(hidden_states, init_states[2])
+        hidden_states = [self.gru_conv_1x1_2(h) for h in hidden_states]
+        hidden_states = [self.g2(h) for h in hidden_states]
+        hidden_states = [self.up_g2(h) for h in hidden_states]
+
+        # Layer 2.
+        hidden_states = self.convGRU3(hidden_states, init_states[1])
+        hidden_states = [self.gru_conv_1x1_3(h) for h in hidden_states]
+        hidden_states = [self.g3(h) for h in hidden_states]
+        hidden_states = [self.up_g3(h) for h in hidden_states]
+
+        # Layer 1 (top-most).
+        hidden_states = self.convGRU4(hidden_states, init_states[0])
+        hidden_states = [self.gru_conv_1x1_4(h) for h in hidden_states]
+        hidden_states = [self.g4(h) for h in hidden_states]
+        hidden_states = [self.up_g4(h) for h in hidden_states]
+
+        # Output layer.
+        hidden_states = [F.relu(self.bn(h)) for h in hidden_states]
+        hidden_states = [self.conv_1x1(h) for h in hidden_states]
+        hidden_states = [self.depth2space(h) for h in hidden_states]
+
+        # Convert forecasts to a torch Tensor
+        forecasts = torch.stack(hidden_states, dim=1)
+        return forecasts
+
+
+class Generator(torch.nn.Module):
+    """Generator class."""
+
+    def __init__(
+        self,
+        conditioning_stack: torch.nn.Module,
+        latent_stack: torch.nn.Module,
+        sampler: torch.nn.Module,
+    ):
+        """
+        Wrap the three parts of the generator for simpler calling.
+
+        Args:
+            conditioning_stack: (torch.nn.Module)
+            latent_stack: (torch.nn.Module)
+            sampler: Combines the conditioning information and latent information (torch.nn.Module)
+        """
+        super().__init__()
+        self.conditioning_stack = conditioning_stack
+        self.latent_stack = latent_stack
+        self.sampler = sampler
+
+    def forward(self, x: torch.Tensor):
+        """Apply a forward pass on the tensor."""
+        conditioning_states = self.conditioning_stack(x)
+        latent_dim = self.latent_stack(x)
+        x = self.sampler(conditioning_states, latent_dim)
+        return x

model/dgmr_official/layers/Attention.py

@@ -0,0 +1,85 @@
+"""Attention module."""
+
+import einops
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+
+def attention_einsum(q, k, v):
+    """Apply the attention operator to tensors of shape [h, w, c]."""
+    # Reshape 3D tensors to 2D tensor with first dimension L = h x w.
+    k = einops.rearrange(k, "h w c -> (h w) c")  # [h, w, c] -> [L, c]
+    v = einops.rearrange(v, "h w c -> (h w) c")  # [h, w, c] -> [L, c]
+
+    # Einstein summation corresponding to the query * key operation.
+    beta = F.softmax(torch.einsum("hwc, Lc->hwL", q, k), dim=-1)
+
+    # Einstein summation corresponding to the attention * value operation.
+    out = torch.einsum("hwL, Lc->hwc", beta, v)
+    return out
+
+
+class AttentionLayer(torch.nn.Module):
+    """Attention layer class."""
+
+    def __init__(
+        self, input_channels: int, output_channels: int, ratio_kq: int = 8, ratio_v: int = 8
+    ):
+        """Initialize the attention layer."""
+        super(AttentionLayer, self).__init__()
+
+        self.ratio_kq = ratio_kq
+        self.ratio_v = ratio_v
+        self.output_channels = output_channels
+        self.input_channels = input_channels
+
+        # Compute query, key and value using 1x1 convolutions.
+        self.query = torch.nn.Conv2d(
+            in_channels=input_channels,
+            out_channels=self.output_channels // self.ratio_kq,
+            kernel_size=(1, 1),
+            padding="valid",
+            bias=False,
+        )
+        self.key = torch.nn.Conv2d(
+            in_channels=input_channels,
+            out_channels=self.output_channels // self.ratio_kq,
+            kernel_size=(1, 1),
+            padding="valid",
+            bias=False,
+        )
+        self.value = torch.nn.Conv2d(
+            in_channels=input_channels,
+            out_channels=self.output_channels // self.ratio_v,
+            kernel_size=(1, 1),
+            padding="valid",
+            bias=False,
+        )
+
+        self.last_conv = torch.nn.Conv2d(
+            in_channels=self.output_channels // 8,
+            out_channels=self.output_channels,
+            kernel_size=(1, 1),
+            padding="valid",
+            bias=False,
+        )
+
+        # Learnable gain parameter
+        self.gamma = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the forward function."""
+        # Compute query, key and value using 1x1 convolutions.
+        query = self.query(x)
+        key = self.key(x)
+        value = self.value(x)
+        # Apply the attention operation.
+        out = []
+        for b in range(x.shape[0]):
+            # Apply to each in batch
+            out.append(attention_einsum(query[b], key[b], value[b]))
+        out = torch.stack(out, dim=0)
+        out = self.gamma * self.last_conv(out)
+        # Residual connection.
+        return out + x

model/dgmr_official/layers/ConvGRU.py

@@ -0,0 +1,111 @@
+"""Implementation of Conv GRU and cell module."""
+
+import torch
+import torch.nn.functional as F
+from torch.nn.utils.parametrizations import spectral_norm
+
+
+class ConvGRUCell(torch.nn.Module):
+    """A ConvGRU implementation."""
+
+    def __init__(
+        self,
+        input_channels: int,
+        output_channels: int,
+        kernel_size: int = 3,
+        sn_eps: float = 0.0001,
+    ):
+        """Conv GRU class.
+
+        Args:
+          input_channels: number of input channels (int)
+          output_channels: number of output channels (int)
+          kernel_size: kernel size of the convolutions. Default: 3.
+          sn_eps: constant for spectral normalization. Default: 1e-4.
+        """
+        super().__init__()
+        self._kernel_size: int = kernel_size
+        self._sn_eps: float = sn_eps
+        self.read_gate_conv = spectral_norm(
+            torch.nn.Conv2d(
+                in_channels=input_channels,
+                out_channels=output_channels,
+                kernel_size=(kernel_size, kernel_size),
+                padding=1,
+            ),
+            eps=sn_eps,
+        )
+        self.update_gate_conv = spectral_norm(
+            torch.nn.Conv2d(
+                in_channels=input_channels,
+                out_channels=output_channels,
+                kernel_size=(kernel_size, kernel_size),
+                padding=1,
+            ),
+            eps=sn_eps,
+        )
+        self.output_conv = spectral_norm(
+            torch.nn.Conv2d(
+                in_channels=input_channels,
+                out_channels=output_channels,
+                kernel_size=(kernel_size, kernel_size),
+                padding=1,
+            ),
+            eps=sn_eps,
+        )
+
+    def forward(self, x, prev_state):
+        """
+        Conv GRU forward, returning the current+new state.
+
+        Args:
+            x: Input tensor
+            prev_state: Previous state
+
+        Returns:
+            New tensor plus the new state
+        """
+        # Concatenate the inputs and previous state along the channel axis.
+        xh = torch.cat([x, prev_state], dim=1)
+
+        # Read gate of the GRU.
+        read_gate = F.sigmoid(self.read_gate_conv(xh))
+
+        # Update gate of the GRU.
+        update_gate = F.sigmoid(self.update_gate_conv(xh))
+
+        # Gate the inputs.
+        gated_input = torch.cat([x, read_gate * prev_state], dim=1)
+
+        # Gate the cell and state / outputs.
+        c = F.relu(self.output_conv(gated_input))
+        out = update_gate * prev_state + (1.0 - update_gate) * c
+        new_state = out
+
+        return out, new_state
+
+
+class ConvGRU(torch.nn.Module):
+    """ConvGRU Cell wrapper to replace tf.static_rnn in TF implementation."""
+
+    def __init__(
+        self,
+        input_channels: int,
+        output_channels: int,
+        kernel_size: int = 3,
+        sn_eps=0.0001,
+    ):
+        """Initialize the convolution layer."""
+        super().__init__()
+        self.cell = ConvGRUCell(input_channels, output_channels, kernel_size, sn_eps)
+
+    def forward(self, x: torch.Tensor, hidden_state=None) -> torch.Tensor:
+        """Apply the forward function on each cell prior to returning it as a stack."""
+        outputs = []
+        for step in range(len(x)):
+            # Compute current timestep
+            output, hidden_state = self.cell(x[step], hidden_state)
+            outputs.append(output)
+        # Stack outputs to return as tensor
+        outputs = torch.stack(outputs, dim=0)
+        return outputs

model/dgmr_official/layers/CoordConv.py

@@ -0,0 +1,81 @@
+"""The coordinate convolution class."""
+
+import torch
+import torch.nn as nn
+
+
+class AddCoords(nn.Module):
+    """argument input tensors with spatial information."""
+
+    def __init__(self, with_r: bool = False):
+        """
+        Initialize the add coordinates class.
+
+        Args:
+            with_r: a condition to check if radical distance should included in the spatial
+            information (bool) default = false
+        """
+        super().__init__()
+        self.with_r: bool = with_r
+
+    def forward(self, input_tensor):
+        """
+        Add spatial information to the input tensor.
+
+        Args:
+            input_tensor: shape(batch, channel, x_dim, y_dim)
+        """
+        batch_size, _, x_dim, y_dim = input_tensor.size()
+
+        xx_channel = torch.arange(x_dim).repeat(1, y_dim, 1)
+        yy_channel = torch.arange(y_dim).repeat(1, x_dim, 1).transpose(1, 2)
+
+        xx_channel = xx_channel.float() / (x_dim - 1)
+        yy_channel = yy_channel.float() / (y_dim - 1)
+
+        xx_channel = xx_channel * 2 - 1
+        yy_channel = yy_channel * 2 - 1
+
+        xx_channel = xx_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3)
+        yy_channel = yy_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3)
+
+        ret = torch.cat(
+            [input_tensor, xx_channel.type_as(input_tensor), yy_channel.type_as(input_tensor)],
+            dim=1,
+        )
+
+        if self.with_r:
+            rr = torch.sqrt(
+                torch.pow(xx_channel.type_as(input_tensor) - 0.5, 2)
+                + torch.pow(yy_channel.type_as(input_tensor) - 0.5, 2)
+            )
+            ret = torch.cat([ret, rr], dim=1)
+
+        return ret
+
+
+class CoordConv(nn.Module):
+    """Coordinate convolution class."""
+
+    def __init__(self, in_channels, out_channels, with_r=False, **kwargs):
+        """
+        Initialize the coordinate convolution.
+
+        Args:
+            in_channels : number of input channels
+            out_channels : number of output channels,
+            with_r : boolean =False,
+            **kwargs : dict[str, Unknown]
+        """
+        super().__init__()
+        self.addcoords = AddCoords(with_r=with_r)
+        in_size = in_channels + 2
+        if with_r:
+            in_size += 1
+        self.conv = nn.Conv2d(in_size, out_channels, **kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply a forward pass on the input tensor."""
+        ret = self.addcoords(x)
+        ret = self.conv(ret)
+        return ret

model/dgmr_official/layers/__init__.py

@@ -0,0 +1,5 @@
+"""Submodules for the layers."""
+
+from .Attention import AttentionLayer
+from .ConvGRU import ConvGRU
+from .CoordConv import CoordConv

model/dgmr_official/layers/utils.py

@@ -0,0 +1,18 @@
+"""Utilities functions for the layer submodules."""
+
+import torch
+
+from . import CoordConv
+
+
+def get_conv_layer(conv_type: str = "standard") -> torch.nn.Module:
+    """Return a conv layer based on the passed in string name."""
+    if conv_type == "standard":
+        conv_layer = torch.nn.Conv2d
+    elif conv_type == "coord":
+        conv_layer = CoordConv
+    elif conv_type == "3d":
+        conv_layer = torch.nn.Conv3d
+    else:
+        raise ValueError(f"{conv_type} is not a recognized Conv method")
+    return conv_layer

model/dgmr_official/losses.py

@@ -0,0 +1,92 @@
+"""Loss functions for DGMR (trimmed from Open Climate Fix ``dgmr`` package).
+
+Only the GAN losses actually used by the DGMR objective are kept; the SSIM
+variants (which depend on ``pytorch_msssim``) are omitted.
+"""
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+
+class GridCellLoss(nn.Module):
+    """Grid Cell Regularizer loss from Skillful Nowcasting, see https://arxiv.org/pdf/2104.00954.pdf."""
+
+    def __init__(self, weight_fn=None, precip_weight_cap=24.0):
+        """
+        Initialize GridCellLoss.
+
+        Args:
+            weight_fn: A function to compute weights for the loss.
+            precip_weight_cap: Custom ceiling value for the weight function.
+        """
+        super().__init__()
+        self.weight_fn = lambda y: weight_fn(y, precip_weight_cap) if weight_fn else None
+
+    def forward(self, generated_images, targets):
+        """
+        Forward function.
+
+        Calculates the grid cell regularizer value, assumes generated images are the mean
+        predictions from 6 calls to the generater (Monte Carlo estimation of the
+        expectations for the latent variable)
+
+        Args:
+            generated_images: Mean generated images from the generator
+            targets: Ground truth future frames
+
+        Returns:
+            Grid Cell Regularizer term
+        """
+        difference = generated_images - targets
+        if self.weight_fn is not None:
+            weights = self.weight_fn(targets)
+            difference = difference * weights
+        difference = difference.norm(p=1)
+        return difference / targets.size(1) * targets.size(3) * targets.size(4)
+
+
+class NowcastingLoss(nn.Module):
+    """Nowcast Loss class."""
+
+    def __init__(self):
+        """Initialize function."""
+        super().__init__()
+
+    def forward(self, x, real_flag):
+        """Apply the relu function to the input tensor."""
+        if real_flag is True:
+            x = -x
+        return F.relu(1.0 + x).mean()
+
+
+def loss_hinge_disc(score_generated, score_real):
+    """Discriminator hinge loss."""
+    l1 = F.relu(1.0 - score_real)
+    loss = torch.mean(l1)
+    l2 = F.relu(1.0 + score_generated)
+    loss += torch.mean(l2)
+    return loss
+
+
+def loss_hinge_gen(score_generated):
+    """Generate hinge loss."""
+    loss = -torch.mean(score_generated)
+    return loss
+
+
+def grid_cell_regularizer(generated_samples, batch_targets):
+    """Grid cell regularizer.
+
+    Args:
+      generated_samples: Tensor of size [n_samples, batch_size, 18, 256, 256, 1].
+      batch_targets: Tensor of size [batch_size, 18, 256, 256, 1].
+
+    Returns:
+      loss: A tensor of shape [batch_size].
+    """
+    gen_mean = torch.mean(generated_samples, dim=0)
+    # TODO Possibly change clip here?
+    weights = torch.clip(batch_targets, 0.0, 24.0)
+    loss = torch.mean(torch.abs(gen_mean - batch_targets) * weights)
+    return loss

scripts/fake_data.py

@@ -0,0 +1,74 @@
+import os
+import h5py
+import numpy as np
+from onescience.utils.YParams import YParams
+
+
+# DGMR 使用 4 帧上下文预测后续雷达帧；T 随模型配置自动变化。
+def get_dims(cfg_model, cfg_data):
+    H, W = map(int, cfg_data.dataset.img_size)
+    if (H, W) != (int(cfg_model.output_shape), int(cfg_model.output_shape)):
+        raise ValueError("dataset.img_size must match model.output_shape")
+    input_steps = int(cfg_model.num_context)
+    output_steps = int(cfg_model.forecast_steps)
+    samples = int(cfg_data.dataloader.batch_size)
+    T = input_steps + output_steps + samples - 1
+    return {
+        "T": T, "H": H, "W": W, "time_step": 1,
+        "input_steps": input_steps, "output_steps": output_steps,
+    }
+
+
+def generate_fake_h5(data_dir, var_names, years, dims):
+    """
+    为每个年份生成一个空 h5 文件。
+    利用 HDF5 chunked 数据集未写入 chunk 即返回 fill_value=0 的特性，
+    文件实际只含元数据，极小，但 shape 与真实数据完全一致。
+    均值/标准差也作为数据集内嵌进每年的 h5，与 era5.py 新版读取方式对应。
+
+    注意：ERA5Datapipe 要求 samples_per_year = T - input_steps - output_steps + 1 >= 1，
+        T 由 num_context、forecast_steps 与 batch_size 自动计算。
+    """
+    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")
+
+
+if __name__ == "__main__":
+    cfg_model = YParams("conf/config.yaml", "model")
+    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
+    if len(atm_vars) != int(cfg_model.input_channels):
+        raise ValueError("channel count must match model.input_channels")
+
+    generate_fake_h5(cfg_datapipe.dataset.data_dir, atm_vars, years, get_dims(cfg_model, cfg_datapipe))
+
+    print("\n✅ Fake datasets generated.")

scripts/inference.py

@@ -0,0 +1,87 @@
+import sys
+from pathlib import Path
+
+# 获取项目根目录（inference.py上级的上级）
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+import torch
+import os
+import glob
+import numpy as np
+import h5py
+from tqdm import tqdm
+from model.dgmr import DGMR
+from onescience.utils.YParams import YParams
+from onescience.datapipes.climate import ERA5Datapipe
+
+
+def get_stats(data_dir, channels):
+    """从新版 h5 中读取变量列表与归一化参数（均值/标准差）"""
+    h5_files = sorted(glob.glob(os.path.join(data_dir, "data", "*.h5")))
+    with h5py.File(h5_files[0], "r") as f:
+        ds = f["fields"]
+        all_variables = [v.decode() if isinstance(v, bytes) else v for v in ds.attrs["variables"]]
+        mu = f["global_means"][:]   # [1, C, 1, 1]
+        std = f["global_stds"][:]
+
+    channel_indices = [all_variables.index(v) for v in channels]
+    means = mu[:, channel_indices, :, :]
+    stds = std[:, channel_indices, :, :]
+    return means, stds
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+
+    ## Model config init
+    config_file_path = os.path.join(current_path, "conf/config.yaml")
+    cfg = YParams(config_file_path, "model")
+
+    ## DataLoader init
+    cfg_data = YParams(config_file_path, "datapipe")
+    means, stds = get_stats(cfg_data.dataset.data_dir, cfg_data.dataset.channels)
+
+    cfg['N_in_channels'] = len(cfg_data.dataset.channels)
+    cfg['N_out_channels'] = len(cfg_data.dataset.channels)
+
+    datapipe = ERA5Datapipe(
+        dataset_dir=cfg_data.dataset.data_dir,
+        used_variables=cfg_data.dataset.channels,
+        used_years=cfg_data.dataset.test_time,
+        distributed=False,
+        input_steps=cfg.num_context,
+        output_steps=cfg.forecast_steps,
+        batch_size=1,
+        num_workers=4,
+    )
+    test_dataloader, _ = datapipe.get_dataloader("test")
+
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location=device, weights_only=False)
+    model = DGMR(
+        forecast_steps=cfg.forecast_steps,
+        num_context=cfg.num_context,
+        input_channels=cfg.input_channels,
+        output_shape=cfg.output_shape,
+        conv_type=cfg.conv_type,
+        latent_channels=cfg.latent_channels,
+        context_channels=cfg.context_channels,
+        generation_steps=cfg.generation_steps,
+        grid_lambda=cfg.grid_lambda,
+        precip_weight_cap=cfg.precip_weight_cap,
+    ).to(device)
+    model.load_state_dict(ckpt["model_state_dict"])
+
+    model.eval()
+    os.makedirs('result/output/', exist_ok=True)
+    print(f"📂 infer results 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, dtype=torch.float32)      # [1, num_context, C, H, W]
+            gen = model.generator(invar).cpu().numpy()           # [1, forecast_steps, C, H, W]
+            for t in range(gen.shape[1]):
+                fname = data[4][cfg.num_context + t][0]          # 该预测帧对应的时刻
+                pred_var = gen[0, t]                             # [C, H, W]
+                pred_var = pred_var * stds + means
+                np.save(f"result/output/{fname}.npy", pred_var)

scripts/result.py

@@ -0,0 +1,203 @@
+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()
+            if pred.ndim == 2:
+                pred = pred[np.newaxis]  # 单通道时 squeeze 会压缩掉通道维，恢复为 [C, H, W]
+
+            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, valid_loss):
+    mask = ~(np.isnan(train_loss) | np.isnan(valid_loss))
+    train_loss = train_loss[mask]
+    valid_loss = valid_loss[mask]
+
+    fig, ax = plt.subplots(figsize=(5, 3.5))
+    colors = {'train': '#2563EB', 'valid': '#EA580C'}
+    epochs = np.arange(1, len(train_loss) + 1)
+
+    ax.plot(epochs, train_loss, color=colors['train'], linewidth=1.5, label='Train')
+    ax.plot(epochs, valid_loss, color=colors['valid'], linewidth=1.5, label='Valid', linestyle='--')
+    min_idx = np.argmin(valid_loss)
+    ax.scatter(epochs[min_idx], valid_loss[min_idx],
+               color=colors['valid'], s=40, zorder=5, edgecolors='white')
+    ax.annotate(f'Best: {valid_loss[min_idx]:.3f}',
+                xy=(epochs[min_idx], valid_loss[min_idx]),
+                xytext=(10, 10), textcoords='offset points', fontsize=8, color=colors['valid'],
+                arrowprops=dict(arrowstyle='-', color=colors['valid'], lw=0.5))
+
+    ax.set(xlabel='Epoch', ylabel='Loss', xlim=(0, len(train_loss) + 1))
+    ax.legend(frameon=False, loc='upper right')
+    ax.grid(True, linestyle='--', alpha=0.3)
+    ax.spines[['top', 'right']].set_visible(False)
+
+    plt.tight_layout()
+    plt.savefig('./result/loss.png', dpi=300, bbox_inches='tight')
+    plt.close()
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    config_file_path = os.path.join(current_path, 'conf/config.yaml')
+    cfg = YParams(config_file_path, 'model')
+    cfg_data = YParams(config_file_path, "datapipe")
+
+    train_loss = np.load('./data/checkpoints/trloss.npy')
+    valid_loss = np.load('./data/checkpoints/valoss.npy')
+    plot_loss(train_loss, valid_loss)
+
+    data_dir = cfg_data.dataset.data_dir
+    total_files, channel_indices, time_step = get_metadata(data_dir, cfg_data.dataset.channels)
+
+    # Load data & Compute RMSE/ACC per channel
+    h5_files = sorted(glob.glob(os.path.join(data_dir, "data", "*.h5")))
+    with h5py.File(h5_files[0], "r") as f:
+        mu = f["global_means"][:]
+    clim_mean = mu[:, channel_indices, :, :]
+    get_result(total_files, channel_indices, time_step, data_dir, clim_mean)
+    show_result()
+
+    ##### 默认绘制 result/output 中第一个预测结果，用户可自行指定日期和变量 #####
+    eg_files = [total_files[0][:-4]]
+    # DGMR 为单通道场，默认绘制该通道（后续帧的时刻可通过 h5 time_step 推算）
+    channel_index = list(range(len(cfg_data.dataset.channels)))
+
+    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()
+        if pred.ndim == 2:
+            pred = pred[np.newaxis]  # 单通道时 squeeze 会压缩掉通道维，恢复为 [C, H, W]
+        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,252 @@
+import sys
+from pathlib import Path
+
+# 获取项目根目录（train.py上级的上级）
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+import torch
+import os
+import numpy as np
+import torch.distributed as dist
+import logging
+import time
+
+from model.dgmr import DGMR
+from model.dgmr_official.losses import loss_hinge_disc, loss_hinge_gen
+from onescience.datapipes.climate import ERA5Datapipe
+from onescience.utils.YParams import YParams
+
+try:
+    from apex import optimizers
+    _FUSED_ADAM = True
+except Exception:
+    _FUSED_ADAM = False
+
+
+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 and torch.cuda.is_available():
+        dist.init_process_group(backend="nccl", init_method="env://")
+        local_rank = int(os.environ["LOCAL_RANK"])
+        world_rank = dist.get_rank()
+    device = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu"
+
+    ## DataLoader init
+    cfg_data = YParams(config_file_path, "datapipe")
+    cfg['N_in_channels'] = len(cfg_data.dataset.channels)
+    cfg['N_out_channels'] = len(cfg_data.dataset.channels)
+    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(),
+        input_steps=cfg.num_context,
+        output_steps=cfg.forecast_steps,
+        batch_size=cfg_data.dataloader.batch_size,
+        num_workers=cfg_data.dataloader.num_workers,
+    )
+    train_dataloader, train_sampler = datapipe.get_dataloader("train")
+    datapipe = ERA5Datapipe(
+        dataset_dir=cfg_data.dataset.data_dir,
+        used_variables=cfg_data.dataset.channels,
+        used_years=cfg_data.dataset.val_time,
+        distributed=dist.is_initialized(),
+        input_steps=cfg.num_context,
+        output_steps=cfg.forecast_steps,
+        batch_size=cfg_data.dataloader.batch_size,
+        num_workers=cfg_data.dataloader.num_workers,
+    )
+    val_dataloader, val_sampler = datapipe.get_dataloader("valid")
+
+    # Model init
+    model = DGMR(
+        forecast_steps=cfg.forecast_steps,
+        num_context=cfg.num_context,
+        input_channels=cfg.input_channels,
+        output_shape=cfg.output_shape,
+        conv_type=cfg.conv_type,
+        latent_channels=cfg.latent_channels,
+        context_channels=cfg.context_channels,
+        generation_steps=cfg.generation_steps,
+        grid_lambda=cfg.grid_lambda,
+        precip_weight_cap=cfg.precip_weight_cap,
+    ).to(device)
+
+    # 生成器与判别器使用独立的 Adam 优化器（论文 lr=1e-4）
+    if _FUSED_ADAM:
+        optimizer_g = optimizers.FusedAdam(model.generator.parameters(), lr=cfg.lr)
+        optimizer_d = optimizers.FusedAdam(model.discriminator.parameters(), lr=cfg.lr_disc)
+    else:
+        optimizer_g = torch.optim.Adam(model.generator.parameters(), lr=cfg.lr)
+        optimizer_d = torch.optim.Adam(model.discriminator.parameters(), lr=cfg.lr_disc)
+    scheduler_g = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_g, factor=0.2, patience=5, mode='min')
+    scheduler_d = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_d, factor=0.2, patience=5, mode='min')
+
+    ## Train process init
+    os.makedirs(cfg.checkpoint_dir, exist_ok=True)
+    train_loss_file = f"{cfg.checkpoint_dir}/trloss.npy"
+    valid_loss_file = f"{cfg.checkpoint_dir}/valoss.npy"
+    best_valid_loss = float("inf")
+    best_loss_epoch = 0
+    train_losses = np.empty((0,), dtype=np.float32)
+    valid_losses = np.empty((0,), dtype=np.float32)
+
+    ## Get model params count
+    if cfg.world_size == 1:
+        total_params = sum(p.numel() for p in model.parameters())
+        print("\n\n")
+        print("-" * 50)
+        print(f"📂 now params is {total_params}, {total_params / 1e6:.2f}M, {total_params / 1e9:.2f}B")
+        print("-" * 50, "\n")
+
+    ## Load model weight if there exist well-trained model
+    if os.path.exists(f"{cfg.checkpoint_dir}/model_bak.pth"):
+        if world_rank == 0:
+            print("\n\n")
+            print("-" * 50)
+            print(f"✅ There has a model weight, load and continue training...")
+            print(f'If you want to train a new model, ensure there is no *.pth file in {cfg.checkpoint_dir}')
+            print("-" * 50, "\n")
+        ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location=device, weights_only=False)
+        model.load_state_dict(ckpt["model_state_dict"])
+        optimizer_g.load_state_dict(ckpt["optimizer_g_state_dict"])
+        optimizer_d.load_state_dict(ckpt["optimizer_d_state_dict"])
+        scheduler_g.load_state_dict(ckpt["scheduler_g_state_dict"])
+        scheduler_d.load_state_dict(ckpt["scheduler_d_state_dict"])
+        best_valid_loss = ckpt["best_valid_loss"]
+        best_loss_epoch = ckpt["best_loss_epoch"]
+        train_losses = np.load(train_loss_file)
+        valid_losses = np.load(valid_loss_file)
+
+    ## Distributed model
+    if dist.is_initialized():
+        model = torch.nn.parallel.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, dtype=torch.float32)      # [B, num_context, C, H, W]
+            outvar = data[1].to(device, dtype=torch.float32)     # [B, forecast_steps, C, H, W]
+            full_real = torch.cat([invar, outvar], dim=1)        # 上下文 + 真实未来帧
+
+            # --- 判别器（hinge loss，生成图像 detach 不传梯度） ---
+            gen_images = model.generator(invar)                  # [B, forecast_steps, C, H, W]
+            full_fake = torch.cat([invar, gen_images], dim=1)
+            score_real = model.discriminator(full_real)
+            score_generated = model.discriminator(full_fake.detach())
+            disc_loss = loss_hinge_disc(score_generated, score_real)
+            optimizer_d.zero_grad()
+            disc_loss.backward()
+            optimizer_d.step()
+
+            # --- 生成器（hinge + 网格单元正则器，MC 采样估计期望） ---
+            score_generated = model.discriminator(full_fake)
+            gen_loss = loss_hinge_gen(score_generated)
+            gen_samples = torch.stack(
+                [model.generator(invar) for _ in range(cfg.generation_steps)], dim=0
+            ).mean(dim=0)                                        # 取 MC 均值作为生成均值图像
+            grid_loss = model.grid_regularizer(gen_samples, outvar)
+            gen_loss = gen_loss + cfg.grid_lambda * grid_loss
+            optimizer_g.zero_grad()
+            gen_loss.backward()
+            optimizer_g.step()
+
+            loss = gen_loss.item() + disc_loss.item()
+            train_loss += loss
+            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'gen_loss:{gen_loss.item(): .04f} disc_loss:{disc_loss.item(): .04f} '
+                            f'loss:{train_loss / (j+1): .04f}')
+
+        train_loss /= len(train_dataloader)
+
+        # 验证只使用网格单元正则器（确定性，单个生成样本）
+        model.eval()
+        valid_loss = 0
+        with torch.no_grad():
+            start_time = time.time()
+            for j, data in enumerate(val_dataloader):
+                invar = data[0].to(device, dtype=torch.float32)
+                outvar = data[1].to(device, dtype=torch.float32)
+                gen_images = model.generator(invar)
+                grid_loss = model.grid_regularizer(gen_images, outvar)
+                loss = grid_loss
+
+                if dist.is_initialized():
+                    loss_tensor = loss.detach().to(device)
+                    dist.all_reduce(loss_tensor)
+                    loss = loss_tensor.item() / cfg.world_size
+                    valid_loss += loss
+                else:
+                    valid_loss += loss.item()
+                if world_rank == 0:
+                    logger.info(f'Valid: Epoch {epoch}-{j+1}/{len(val_dataloader)} '
+                            f'[cost {int((time.time()-start_time) // 60):02}:{int((time.time()-start_time) % 60):02}] '
+                            f'loss:{valid_loss / (j+1): .04f}')
+
+        valid_loss /= len(val_dataloader)
+        is_save_ckp = False
+        if valid_loss < best_valid_loss:
+            best_valid_loss = valid_loss
+            best_loss_epoch = epoch
+            world_rank == 0 and save_checkpoint(model, optimizer_g, optimizer_d, scheduler_g, scheduler_d, best_valid_loss, best_loss_epoch, cfg.checkpoint_dir)
+            is_save_ckp = True
+        scheduler_g.step(valid_loss)
+        scheduler_d.step(valid_loss)
+
+        if world_rank == 0:
+            logger.info(f"Epoch [{epoch + 1}/{cfg.max_epoch}], "
+                        f"Train Loss: {train_loss:.4f}, "
+                        f"Valid Loss: {valid_loss:.4f}, "
+                        f"Best loss at Epoch: {best_loss_epoch + 1}"
+                        + (", saving checkpoint" if is_save_ckp else "")
+                        )
+            train_losses = np.append(train_losses, train_loss)
+            valid_losses = np.append(valid_losses, valid_loss)
+            np.save(train_loss_file, train_losses)
+            np.save(valid_loss_file, valid_losses)
+
+        if epoch - best_loss_epoch > cfg.patience:
+            print(f"Loss has not decrease in {cfg.patience} epochs, stopping training...")
+            exit()
+
+
+def save_checkpoint(model, optimizer_g, optimizer_d, scheduler_g, scheduler_d, 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_g_state_dict": optimizer_g.state_dict(),
+             "optimizer_d_state_dict": optimizer_d.state_dict(),
+             "scheduler_g_state_dict": scheduler_g.state_dict(),
+             "scheduler_d_state_dict": scheduler_d.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()

weight/.gitkeep

@@ -0,0 +1 @@
+