跳转至

DGMR(deep generative models of radar)

暂无

# Download data from Huggingface
mkdir openclimatefix/nimrod-uk-1km/20200718/valid/subsampled_tiles_256_20min_stride
cd openclimatefix/nimrod-uk-1km/20200718/valid/subsampled_tiles_256_20min_stride
git lfs install
git lfs pull --include="seq-24-*-of-00033.tfrecord.gz"

python dgmr.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/dgmr/dgmr_pretrained.pdparams
预训练模型 指标
dgmr_pretrained.pdparams d_loss: 127.041
g_loss: 59.2409
grid_loss: 1.7699

1. 背景简介

短临降水预报是对未来两小时内的降水进行高分辨率预测，支持了许多依赖于天气决策的实际社会经济需求。最先进的运行即时预报方法通常利用基于雷达的风估计对降水场进行平流，但往往难以捕捉重要的非线性事件，如对流的发生。最近引入的深度学习方法利用雷达直接预测未来的降雨率，摆脱了物理约束。虽然它们能够准确预测低强度降雨，但由于缺乏约束，在更长的前瞻时间内产生模糊的即时预报，导致对中到大雨事件的性能较差，其运行实用性受到限制。为了解决这些挑战，我们提出了一种用于基于雷达的降水概率即时预报的深度生成模型。我们的模型在范围为1536 km × 1280 km的区域内，能够在5-90分钟的前瞻时间内生成逼真且时空一致的预测。通过英国气象局的五十多位专家预报员进行的系统评估显示，我们的生成模型在88%的情况下在准确性和实用性方面排名第一，超过了两种竞争方法，表明了其在决策价值和向现实世界专家提供物理洞察力方面的能力。在定量验证方面，这些即时预报具有良好的技能而无需进行模糊处理。我们展示了生成式即时预报可以提供改进预报价值并支持运行实用性的概率性预测，在分辨率和前瞻时间方面，替代方法存在困难的情况下尤其如此。

短临降水预报在工程和科学领域具有多方面的重要性，主要体现在以下几个方面：

  • 社会影响： 短临降水预报对各行各业的实际决策都有着直接的社会影响。例如，农业、水资源管理、城市防汛和交通运输等领域都需要准确的降水预报来做出相应的应对措施，以减少可能的损失和风险。
  • 安全保障： 短临降水预报对于保障公众安全至关重要。例如，预警系统可以根据短临降水预报提前通知人们可能发生的暴雨、洪涝、泥石流等灾害，从而及时采取避险措施，减少人员伤亡和财产损失。
  • 生态环境保护： 对降水的准确预报有助于生态环境的保护和管理。例如，预测降水量可以帮助决策者及时调整水利工程的运行，保障生态系统的健康运行，并为植被生长提供必要的水资源。
  • 科学研究： 短临降水预报也是气象科学研究的重要组成部分。通过对降水过程的研究和预测，可以更好地理解大气环境中的水循环过程，为气象学、气候学等相关领域的研究提供重要数据和支撑。
  • 工程规划和设计： 在城市规划、土木工程、农业灌溉等领域，准确的短临降水预报对工程的规划和设计至关重要。例如，在城市排水系统设计中，需要考虑未来短时间内可能发生的降水情况，以保证排水系统的正常运行和城市的防洪能力。

总的来说，短临降水预报在工程和科学领域的重要性体现在保障社会安全、促进科学研究、支持生态环境保护和推动工程发展等多个方面，对于各行业的发展和社会的可持续发展都具有重要意义。

2. 模型原理

2.1 模型结构

DGMR是在条件生成对抗网络的算法框架中构建的。以过去的雷达数据为基础，对未来的雷达做出详细和可信的预测。也就是说在给定的时间点 \(T\)，使用基于雷达的地表降水估计值 \(X_T\)，基于过去 \(M\) 个雷达场预测未来 \(N\) 个雷达场。即：

\[ P\left(X_{M+1: M+N} \mid X_{1: M}\right)=\int P\left(X_{M+1: M+N} \mid \mathrm{Z}, X_{1: M}, \boldsymbol{\theta}\right) P\left(\mathrm{Z} \mid X_{1: M}\right) d \mathrm{Z}. \]

其中 \(Z\) 为随机向量，\(\theta\) 为生成模型的参数。

  • 等式左边是条件概率，给定过去\(M\)个时刻的雷达降水，预报之后\(N\)个时刻的雷达降水。
  • 右边则将概率写为集合预报的积分形式：
  • 给定随机抽样 \(Z\) 和生成网络参数 \(\theta\)，在过去 \(M\) 个时刻的雷达降水约束下预报之后 \(N\) 个时刻的雷达降水；
  • 计算随机抽样 \(Z\) 在过去 \(M\) 个时刻的雷达降水约束下的条件概率；
  • 两者相乘，为该结果的出现概率，积分后得到多次抽样下的集合预报。

对随机向量 \(Z\) 的积分确保了模型产生的预测具有空间相关性。DGMR 专门用于降水预测问题。四个连续的雷达观测数据（前20分钟）被用作生成器的输入，该生成器允许对未来降水的多个实现进行抽样，每个实现包含18帧（90分钟）。模型架构示意图如图所示。

domain_chip.pdf

模型架构示意图。

DGMR 是一个使用两个判别器和一个附加正则化项进行训练的生成器。下图显示了生成模型和判别器的详细示意图：

domain_chip.pdf

a、生成器架构。b，生成器的时间鉴别器架构（左上）、空间鉴别器（左中）和潜在条件堆栈（左下）。右侧是 G 块（上）、D 和 3D 块（中）以及 L 块（右）的架构。

2.2 目标函数

生成器通过两个鉴别器的损失和一个网格单元正则化项（记为 \(\mathcal{L}_R(\theta)\) ）进行训练。空间鉴别器 \(D\phi\) 具有参数 \(\phi\)，时间鉴别器 \(T_\psi\) 具有参数 \(\psi\)，生成器 \(G_\theta\) 具有参数 \(\theta\)。我们使用符号 \(\{X ; G\}\) 表示两个字段的串联。最大化的生成器损失如下：

\[ \begin{gathered} \mathcal{L}_G(\theta)=\mathbb{E}_{X_{1: M+N}}\left[\mathbb{E}_Z\left[D\left(G_\theta\left(Z ; X_{1: M}\right)\right)+T\left(\left\{X_{1: M} ; G_\theta\left(Z ; X_{1: M}\right)\right\}\right)\right]-\lambda \mathcal{L}_R(\theta)\right] ; \\ \mathcal{L}_R(\theta)=\frac{1}{H W N}\left\|\left(\mathbb{E}_Z\left[G_\theta\left(Z ; X_{1: M}\right)\right]-X_{M+1: M+N}\right) \odot w\left(X_{M+1: M+N}\right)\right\|_1 . \end{gathered} \]

我们在上面公式中对潜变量 \(\mathrm{Z}\) 的期望使用 Carlo 估计。这些估计是使用每个输入 \(X_{1: M}\) 的六个样本计算的，其中包括 \(M=4\) 个雷达观测数据。网格单元正则化项确保平均预测保持接近真实值，并在高度 \(H\)、宽度 \(W\) 和提前时间 \(N\) 轴上对所有网格单元进行平均。它通过函数 \(w(y)=\max (y+1,24)\) 加权至更高的降雨目标，该函数对输入向量进行逐元素操作，并在 24 处截断以提高对雷达中异常大值的鲁棒性。GAN 空间鉴别器损失 \(\mathcal{L}_D(\phi)\) 和时间鉴别器损失 \(\mathcal{L}_T(\psi)\) 分别相对于参数 \(\phi\) 和 \(\psi\) 最小化。鉴别器损失采用铰链损失公式：

\[ \begin{aligned} & \mathcal{L}_D(\phi)=\mathbb{E}_{X_{1: M+N}, Z}\left[\operatorname{ReLU}\left(1-D_\phi\left(X_{M+1: M+N}\right)\right)+\operatorname{ReLU}\left(1+D_\phi\left(G\left(Z ; X_{1: M}\right)\right)\right)\right], \\ & \mathcal{L}_T(\psi)=\mathbb{E}_{X_{1: M+N}, Z}\left[\operatorname{ReLU}\left(1-T_\psi\left(X_{1: M+N}\right)\right)+\operatorname{ReLU}\left(1+T_\psi\left(\left\{X_{1: M} ; G\left(Z ; X_{1: M}\right)\right\}\right)\right)\right], \end{aligned} \]

其中 \(\operatorname{ReLU} = \max(0,x)\). 更多详细的理论推导请参考 Skillful Precipitation Nowcasting using Deep Generative Models of Radar。

3. 问题求解

接下来开始讲解如何将该问题一步一步地转化为 PaddleScience 代码，用 DGMR 来预测短临降水。为了快速理解 PaddleScience，接下来仅对模型构建、约束构建等关键步骤进行阐述，而其余细节请参考API文档。

3.1 数据集介绍

为了训练和评估英国的临近预报模型，DGMR 使用了英国气象局 RadarNet4 网络中的雷达复合数据。使用 2016 年 1 月 1 日至 2019 年 12 月 31 日期间每五分钟收集一次的雷达数据。我们使用以下数据分割进行模型开发。将 2016 年至 2018 年每月第一天的字段分配给验证集。2016 年至 2018 年的所有其他日期都分配给训练集。最后，使用2019年的数据作为测试集，防止数据泄露和分布泛化测试。开源的英国训练数据集已镜像到 HuggingFace 数据集，用户可以自行下载使用。

在该模型的PaddleScience 代码中，我们可以调用ppsci.data.dataset.DGMRDataset 来加载数据集，代码如下：

199
200
201
# load evaluate data
dataset = ppsci.data.dataset.DGMRDataset(**cfg.DATASET)
val_loader = paddle.io.DataLoader(dataset, batch_size=4)

3.2 模型构建

在 DGMR 模型中，输入过去四个雷达场数据，对 18 个未来雷达场（接下来的 90 分钟）进行预测。DGMR 网络可以表示为 \(X_{1:4}\) 到输出 \(X_{5:22}\) 的映射函数 \(f\)，即：

\[ X_{5:22} = f(X_{1:4}),\\ \]

上式中 \(f\) 代表 DGMR 模型。我们定义 PaddleScience 内置的 DGMR 模型类，并调用，PaddleScience 代码表示如下

197
198
# set model
model = ppsci.arch.DGMR(**cfg.MODEL)

这样我们就实例化出了一个 DGMR 模型，模型参数设置相关内容请参考文献 Skillful Precipitation Nowcasting using Deep Generative Models of Radar。

3.3 模型评估、可视化

模型构建和加载数据后，将上述实例化的对象按顺序传递给 ppsci.solver.Solver，然后启动评估和可视化。首先我们初始化模型：

202
203
204
205
206
207
# initialize solver
solver = ppsci.solver.Solver(
    model,
    pretrained_model_path=cfg.EVAL.pretrained_model_path,
)
solver.model.eval()

然后自定义评估方式和损失函数，代码如下：

 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def validation(
    cfg: DictConfig,
    solver: ppsci.solver.Solver,
    batch: Tuple[Dict[str, paddle.Tensor], ...],
):
    """
    validation step.

    Args:
        cfg (DictConfig): Configuration object.
        solver (ppsci.solver.Solver): Solver object containing the model and related components.
        batch (Tuple[Dict[str, paddle.Tensor], ...]): Input batch consisting of images and corresponding future images.

    Returns:
        discriminator_loss: Loss incurred by the discriminator.
        generator_loss: Loss incurred by the generator.
        grid_cell_reg: Regularization term to encourage smooth transitions.
    """
    images, future_images = batch
    images_value = images[cfg.DATASET.input_keys[0]]
    future_images_value = future_images[cfg.DATASET.label_keys[0]]
    # Two discriminator steps per generator step
    for _ in range(2):
        predictions = solver.predict(images)
        predictions_value = predictions[cfg.MODEL.output_keys[0]]
        generated_sequence = paddle.concat(x=[images_value, predictions_value], axis=1)
        real_sequence = paddle.concat(x=[images_value, future_images_value], axis=1)
        concatenated_inputs = paddle.concat(
            x=[real_sequence, generated_sequence], axis=0
        )
        concatenated_outputs = solver.model.discriminator(concatenated_inputs)
        score_real, score_generated = paddle.split(
            x=concatenated_outputs,
            num_or_sections=[real_sequence.shape[0], generated_sequence.shape[0]],
            axis=0,
        )
        score_real_spatial, score_real_temporal = paddle.split(
            x=score_real, num_or_sections=score_real.shape[1], axis=1
        )
        score_generated_spatial, score_generated_temporal = paddle.split(
            x=score_generated, num_or_sections=score_generated.shape[1], axis=1
        )
        discriminator_loss = _loss_hinge_disc(
            score_generated_spatial, score_real_spatial
        ) + _loss_hinge_disc(score_generated_temporal, score_real_temporal)

    predictions_value = [
        solver.predict(images)[cfg.MODEL.output_keys[0]] for _ in range(6)
    ]
    grid_cell_reg = _grid_cell_regularizer(
        paddle.stack(x=predictions_value, axis=0), future_images_value
    )
    generated_sequence = [
        paddle.concat(x=[images_value, x], axis=1) for x in predictions_value
    ]
    real_sequence = paddle.concat(x=[images_value, future_images_value], axis=1)
    generated_scores = []
    for g_seq in generated_sequence:
        concatenated_inputs = paddle.concat(x=[real_sequence, g_seq], axis=0)
        concatenated_outputs = solver.model.discriminator(concatenated_inputs)
        score_real, score_generated = paddle.split(
            x=concatenated_outputs,
            num_or_sections=[real_sequence.shape[0], g_seq.shape[0]],
            axis=0,
        )
        generated_scores.append(score_generated)
    generator_disc_loss = _loss_hinge_gen(paddle.concat(x=generated_scores, axis=0))
    generator_loss = generator_disc_loss + 20 * grid_cell_reg

    return discriminator_loss, generator_loss, grid_cell_reg


def _loss_hinge_disc(score_generated, score_real):
    """Discriminator hinge loss."""
    l1 = nn.functional.relu(x=1.0 - score_real)
    loss = paddle.mean(x=l1)
    l2 = nn.functional.relu(x=1.0 + score_generated)
    loss += paddle.mean(x=l2)
    return loss


def _loss_hinge_gen(score_generated):
    """Generator hinge loss."""
    loss = -paddle.mean(x=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 = paddle.mean(x=generated_samples, axis=0)
    weights = paddle.clip(x=batch_targets, min=0.0, max=24.0)
    loss = paddle.mean(x=paddle.abs(x=gen_mean - batch_targets) * weights)
    return loss

最后对数据中的每个 batch 进行遍历评估，同时对预测结果进行可视化。

209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
# evaluate pretrained model
d_loss = []
g_loss = []
grid_loss = []
for batch_idx, batch in enumerate(val_loader):
    with paddle.no_grad():
        out_dict = validation(cfg, solver, batch)

        # visualize
        images = batch[0][cfg.DATASET.input_keys[0]]
        future_images = batch[1][cfg.DATASET.label_keys[0]]
        generated_images = solver.predict(batch[0])[cfg.MODEL.output_keys[0]]
        if batch_idx % 50 == 0:
            logger.message(f"Saving plot of image frame to {cfg.output_dir}")
            visualize(
                cfg.output_dir, images, future_images, generated_images, batch_idx
            )

    d_loss.append(out_dict[0])
    g_loss.append(out_dict[1])
    grid_loss.append(out_dict[2])
logger.message(f"d_loss: {np.array(d_loss).mean()}")
logger.message(f"g_loss: {np.array(g_loss).mean()}")
logger.message(f"grid_loss: {np.array(grid_loss).mean()}")

4. 完整代码

dgmr.py
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Reference: https://github.com/openclimatefix/skillful_nowcasting
"""
from os import path as osp
from typing import Dict
from typing import Tuple

import hydra
import matplotlib.pyplot as plt
import numpy as np
import paddle
import paddle.nn as nn
from omegaconf import DictConfig

import ppsci
from ppsci.utils import logger


def visualize(
    output_dir: str,
    x: paddle.Tensor,
    y: paddle.Tensor,
    y_hat: paddle.Tensor,
    batch_idx: int,
) -> None:
    """
    Visualizes input, target, and generated images and saves them to the output directory.

    Args:
        output_dir (str): Directory to save the visualization images.
        x (paddle.Tensor): Input images tensor.
        y (paddle.Tensor): Target images tensor.
        y_hat (paddle.Tensor): Generated images tensor.
        batch_idx (int): Batch index.

    Returns:
        None
    """
    images = x[0]
    future_images = y[0]
    generated_images = y_hat[0]
    fig, axes = plt.subplots(2, 2)
    for i, ax in enumerate(axes.flat):
        alpha = images[i][0].numpy()
        alpha[alpha < 1] = 0
        alpha[alpha > 1] = 1
        ax.imshow(images[i].transpose([1, 2, 0]).numpy(), alpha=alpha, cmap="viridis")
    plt.subplots_adjust(hspace=0.1, wspace=0.1)
    plt.savefig(osp.join(output_dir, f"Input_Image_Stack_Frame_{batch_idx}.png"))
    fig, axes = plt.subplots(3, 3)
    for i, ax in enumerate(axes.flat):
        alpha = future_images[i][0].numpy()
        alpha[alpha < 1] = 0
        alpha[alpha > 1] = 1
        ax.imshow(
            future_images[i].transpose([1, 2, 0]).numpy(), alpha=alpha, cmap="viridis"
        )
    plt.subplots_adjust(hspace=0.1, wspace=0.1)
    plt.savefig(osp.join(output_dir, f"Target_Image_Frame_{batch_idx}.png"))
    fig, axes = plt.subplots(3, 3)
    for i, ax in enumerate(axes.flat):
        alpha = generated_images[i][0].numpy()
        alpha[alpha < 1] = 0
        alpha[alpha > 1] = 1
        ax.imshow(
            generated_images[i].transpose([1, 2, 0]).numpy(),
            alpha=alpha,
            cmap="viridis",
        )
    plt.subplots_adjust(hspace=0.1, wspace=0.1)
    plt.savefig(osp.join(output_dir, f"Generated_Image_Frame_{batch_idx}.png"))
    plt.close()


def validation(
    cfg: DictConfig,
    solver: ppsci.solver.Solver,
    batch: Tuple[Dict[str, paddle.Tensor], ...],
):
    """
    validation step.

    Args:
        cfg (DictConfig): Configuration object.
        solver (ppsci.solver.Solver): Solver object containing the model and related components.
        batch (Tuple[Dict[str, paddle.Tensor], ...]): Input batch consisting of images and corresponding future images.

    Returns:
        discriminator_loss: Loss incurred by the discriminator.
        generator_loss: Loss incurred by the generator.
        grid_cell_reg: Regularization term to encourage smooth transitions.
    """
    images, future_images = batch
    images_value = images[cfg.DATASET.input_keys[0]]
    future_images_value = future_images[cfg.DATASET.label_keys[0]]
    # Two discriminator steps per generator step
    for _ in range(2):
        predictions = solver.predict(images)
        predictions_value = predictions[cfg.MODEL.output_keys[0]]
        generated_sequence = paddle.concat(x=[images_value, predictions_value], axis=1)
        real_sequence = paddle.concat(x=[images_value, future_images_value], axis=1)
        concatenated_inputs = paddle.concat(
            x=[real_sequence, generated_sequence], axis=0
        )
        concatenated_outputs = solver.model.discriminator(concatenated_inputs)
        score_real, score_generated = paddle.split(
            x=concatenated_outputs,
            num_or_sections=[real_sequence.shape[0], generated_sequence.shape[0]],
            axis=0,
        )
        score_real_spatial, score_real_temporal = paddle.split(
            x=score_real, num_or_sections=score_real.shape[1], axis=1
        )
        score_generated_spatial, score_generated_temporal = paddle.split(
            x=score_generated, num_or_sections=score_generated.shape[1], axis=1
        )
        discriminator_loss = _loss_hinge_disc(
            score_generated_spatial, score_real_spatial
        ) + _loss_hinge_disc(score_generated_temporal, score_real_temporal)

    predictions_value = [
        solver.predict(images)[cfg.MODEL.output_keys[0]] for _ in range(6)
    ]
    grid_cell_reg = _grid_cell_regularizer(
        paddle.stack(x=predictions_value, axis=0), future_images_value
    )
    generated_sequence = [
        paddle.concat(x=[images_value, x], axis=1) for x in predictions_value
    ]
    real_sequence = paddle.concat(x=[images_value, future_images_value], axis=1)
    generated_scores = []
    for g_seq in generated_sequence:
        concatenated_inputs = paddle.concat(x=[real_sequence, g_seq], axis=0)
        concatenated_outputs = solver.model.discriminator(concatenated_inputs)
        score_real, score_generated = paddle.split(
            x=concatenated_outputs,
            num_or_sections=[real_sequence.shape[0], g_seq.shape[0]],
            axis=0,
        )
        generated_scores.append(score_generated)
    generator_disc_loss = _loss_hinge_gen(paddle.concat(x=generated_scores, axis=0))
    generator_loss = generator_disc_loss + 20 * grid_cell_reg

    return discriminator_loss, generator_loss, grid_cell_reg


def _loss_hinge_disc(score_generated, score_real):
    """Discriminator hinge loss."""
    l1 = nn.functional.relu(x=1.0 - score_real)
    loss = paddle.mean(x=l1)
    l2 = nn.functional.relu(x=1.0 + score_generated)
    loss += paddle.mean(x=l2)
    return loss


def _loss_hinge_gen(score_generated):
    """Generator hinge loss."""
    loss = -paddle.mean(x=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 = paddle.mean(x=generated_samples, axis=0)
    weights = paddle.clip(x=batch_targets, min=0.0, max=24.0)
    loss = paddle.mean(x=paddle.abs(x=gen_mean - batch_targets) * weights)
    return loss


def train(cfg: DictConfig):
    raise NotImplementedError("Training of DGMR is not supported now.")


def evaluate(cfg: DictConfig):
    # set model
    model = ppsci.arch.DGMR(**cfg.MODEL)
    # load evaluate data
    dataset = ppsci.data.dataset.DGMRDataset(**cfg.DATASET)
    val_loader = paddle.io.DataLoader(dataset, batch_size=4)
    # initialize solver
    solver = ppsci.solver.Solver(
        model,
        pretrained_model_path=cfg.EVAL.pretrained_model_path,
    )
    solver.model.eval()

    # evaluate pretrained model
    d_loss = []
    g_loss = []
    grid_loss = []
    for batch_idx, batch in enumerate(val_loader):
        with paddle.no_grad():
            out_dict = validation(cfg, solver, batch)

            # visualize
            images = batch[0][cfg.DATASET.input_keys[0]]
            future_images = batch[1][cfg.DATASET.label_keys[0]]
            generated_images = solver.predict(batch[0])[cfg.MODEL.output_keys[0]]
            if batch_idx % 50 == 0:
                logger.message(f"Saving plot of image frame to {cfg.output_dir}")
                visualize(
                    cfg.output_dir, images, future_images, generated_images, batch_idx
                )

        d_loss.append(out_dict[0])
        g_loss.append(out_dict[1])
        grid_loss.append(out_dict[2])
    logger.message(f"d_loss: {np.array(d_loss).mean()}")
    logger.message(f"g_loss: {np.array(g_loss).mean()}")
    logger.message(f"grid_loss: {np.array(grid_loss).mean()}")


@hydra.main(version_base=None, config_path="./conf", config_name="dgmr.yaml")
def main(cfg: DictConfig):
    if cfg.mode == "train":
        train(cfg)
    elif cfg.mode == "eval":
        evaluate(cfg)
    else:
        raise ValueError(f"cfg.mode should in ['train', 'eval'], but got '{cfg.mode}'")


if __name__ == "__main__":
    main()

5. 结果展示

如图所示为分别在 \(T+5, T+10, \cdots, T+45\) 分钟的气象降水预测，与真实的降水情况相比可以看出，该模型可以给出比较好的短临降水预测。

chip.png

模型预测的降水情况

chip.png

真实降水情况

6. 参考资料

参考文献： Skillful Precipitation Nowcasting using Deep Generative Models of Radar

参考代码： DGMR