原文地址

雷达回波外推是一个还算处女的方向,本质上来说这就是一个视频序列预测视频序列的问题,比如这里的使用已知的 5 帧预测未来的 20 帧雷达回波序列(通常每 6 分钟采集一次)。所以这个问题是可以使用现有的一些处理视频的方法来解决的,例如 RNN 和 Conv3D 等能处理时空信息的结构。

目前这个领域中文的一些资料很少有人介绍和总结。趁着还有些记忆,写点东西记录一下。在这里我会介绍雷达回波外推这个问题、HKO-7 数据集以及施行健的两篇 nips 文章,分别是 ConvLSTM 和 TrajGRU。

我用 pytorch 实现了这两篇文章,并在 HKO-7 数据集上取得了和作者近似的结果。原始代码在 HKO-7,使用 mxnet 实现,其中有 ConvGRU 和 TrajGRU,没有 ConvLSTM。作者的代码实现了整个数据加载、训练和测试流程,封装的也非常好,我在实现时也使用了一部分作者的代码,我认为我的实现要更容易理解一些。

代码挂在了 Github。如果你觉得这篇文章或代码对你有一定帮助,欢迎点赞这篇文章或 star 这个 repo。

在开始之前,可以通过 Machine Learning for Spatiotemporal Sequence Forecasting and Its Application to Nowcasting香港科技大学施行健:深度学习用于短临降雨预报的一个基准和一个新模型 | 分享总结 对雷达回波外推问题有一个比较清晰的认识。

数据集

雷达回波数据主要是通过天气雷达往四周发射电磁波并通过云层的反射得到,国内用的大多是多普勒雷达,下面是雷达的扫描示意图,天气雷达的覆盖区域是一个圆形区域:

因为是通过反射得到的,所以雷达一般部署在空旷地区。如果四周出现了如建筑物、灰尘等物体,会在雷达回波中出现明显的噪声,使模型难以收敛。

现在能找到的开源数据集主要有三个,其中两个来自于天池比赛,分别是 Tianchi CIKM AnalytiCup 2017Tianchi IEEE ICDM 2018 全球气象AI挑战赛。这两个比赛最终目的不相同,但是数据集的含义是一样的,降水量等价于雷达反射的 dBZ。这两个比赛一些分享的经验,有一定参考意义,最主要的是 Tianchi IEEE ICDM 2018 全球气象AI挑战赛 这个比赛的主题才是雷达回波外推,但是没有找到开源的比较好的解决方案。

如果希望使用到私有的数据集可以向本地的气象台询问,或者在 中国气象数据网 下载多普勒雷达基数据。后者我尝试过实名注册,个人用户每日限额 2M,想真正用是不可能的,只能采用前面介绍的方法,或者注册企业单位(可能收费)。

最后则是 HKO-7 Dataset,这个数据集是施行健在 Deep Learning for Precipitation Nowcasting: A Benchmark and A New Model 针对雷达回波外推这个领域没有统一的 benchmark 而提出的一个数据集,来源于香港天文台。因为不开放所以需要申请,填张表就 OK,学生的话要找导师去申请,然后用工作邮箱发过去。

这个数据集覆盖了香港周边的雷达回波数据,图片大小是 $480\times 480$ 的灰度图,涵盖了 2009 到 2015 的 下雨天 的数据,812 天用于训练、50 天用于验证、131 天用于测试。更详细信息可以去阅读论文。

首先数据集的转换关系是这样的:dBZ 是雷达的基本反射率(雷达扫出来的结果),然后通过 $\left[ 255 \times \frac { \mathrm { dBZ } + 10 } { 70 } + 0.5 \right]$ 并裁剪到 $[0, 255]$ 之间转换成普通的 png 图片。这种转换对于最后的训练和测试没有多大的影响,只是个线性转换,在深度学习中这个范围才是更常见也更容易解决问题的。

最后还有一种转换关系则是通过 Z-R 关系在 dBZ 和降水值 R(单位是 mm/h)之间转换:$\mathrm { dBZ } = 10 \log a + 10 b \log R, a = 58.53,b = 1.56$,$a$,$b$ 通过线性回归计算。

作者统计了全部数据集的降水值分布:

由于降水值分布不均匀,并且大雨对现实生活的影响更大,所以 使用了一个加权损失函数 进行训练(这篇文章采用的是 MSE 和 MAE 之和)。

$$ w ( x ) = \left\{ \begin{array} { l l } { 1 , } & { x < 2 } \\ { 2 , } & { 2 \leq x < 5 } \\ { 5 , } & { 5 \leq x < 10 } \\ { 10 , } & { 10 \leq x < 30 } \\ { 30 , } & { x \geq 30 } \end{array} \right. $$

在数据预处理上,除了以上介绍的,由于原始雷达回波数据还存在大量噪声,作者采用马氏距离去除噪声点,得到一个 mask 之后,只有正常数据才会参与训练和评估。因此整个数据集分成了原始 png 和对应的 mask 两个部分。

具体过程是:计算每一张图片的均值 $ \hat { \mu } = \frac { \sum _ { i = 1 } ^ { N } x _ { i } } { N } $ 和协方差矩阵 $\hat { \mathbf { S } } = \frac { \sum _ { i = 1 } ^ { N } \left( x _ { i } - \mu \right) \left( x _ { i } - \mu \right) ^ { T } } { N - 1 }$,然后计算出马氏距离: $$D _ { M } ( x ) =\sqrt { ( x - \hat { \mu } ) ^ { T } \hat { \mathbf { S } } ^ { \dagger } ( x - \hat { \mu } ) ^ { 4 } }$$

去除掉马氏距离大于 $\mu + 3\sigma$ 和像素值在 $[0, 71]$ 之间的数据,最后得到一个掩膜就是正常的数据。第一步是为了去除一些离散点,因为存在噪声;第二步则是去除雷达监测范围以外的区域,因为雷达扫描区域是一个圆形的。最后的效果则是:

作者采用的是用 5 帧预测未来的 20 帧。在训练的时候计算量非常大,我在实现过程中一颗 1070 的 batch_size 只能为 2,迭代 100000 次需要三天。由于这个数据量是完全足够的,所以不需要什么数据扩充,只做了一个归一化到 1 以内。

我采用的是作者实现的数据加载的代码,训练时从数据集随机抽取连续的 25 帧(已知 5 帧,预测 20 帧),而评估时则是按 5 的窗口进行滑动,直到结束。这种方式在实际用的时候最好还是上多显卡进行训练,并且加载数据时采用 prefetcher 提前加载数据集,否则在训练的时候大部分时间都会用在数据加载上,浪费很多时间。关于如何加速训练可以参考 如何给你PyTorch里的Dataloader打鸡血

雷达回波外推的评价标准主要使用 HSS、CSI、POD 等指标,参考 评价标准,这几个指标和机器学习中的 precision 和 recall 类似,越高越好。其他的还有 MSE、MAE,施行健采用的 B-MSE 和 B-MAE 其实就是加权的 MSE、MAE,这两个指标越低越好。

ConvLSTM

ConvLSTM 来源于 Convolutional LSTM Network: A Machine Learning Approach for Precipitation Nowcasting

传统的 FC-LSTM 将数据展开成一维进行预测,雷达回波数据存在大量冗余信息,FC-LSTM 无法处理,而且 FC-LSTM 只能提取时间序列信息,无法提取空间信息。

ConvLSTM 使用卷积替换 FC-LSTM,利用卷积操作提取空间信息,FC-LSTM 和 ConvLSTM 的区别主要在输入 $\mathcal { X } _ { 1 } , \ldots , \mathcal { X } _ { t }$,候选态 $\mathcal { C } _ { 1 } , \ldots , \mathcal { C } _ { t }$,隐藏态 $\mathcal { H } _ { 1 } , \ldots , \mathcal { H } _ { t }$ 和三个门 $i _ { t } , f _ { t } , o _ { t }$ 都是三维张量,其中最后两个维度是 spatial 空间维度。如下图:

对比两个结构的公式,LSTM 的公式: $$ \begin{equation} \begin{aligned} i _ { t } &= \sigma \left( W _ { x i } \mathrm { X } _ { t } + W _ { h i } \mathrm { H } _ { t - 1 } + W _ { c i } \circ \mathrm { C } _ { t - 1 } + b _ { i } \right)\\ f _ { t } &= \sigma \left( W _ { x f } \mathrm { X } _ { t } + W _ { h f } \mathrm { H } _ { t - 1 } + W _ { c f } \circ \mathrm { C } _ { i - 1 } + b _ { f } \right)\\ C _ { t } &= f _ { t } \circ C _ { t - 1 } + i _ { t } \circ \tanh \left( W _ { x c } X _ { t } + W _ { h c } H _ { t - 1 } + b _ { c } \right)\\ o _ { t } &= \sigma \left( W _ { x o } \mathrm { X } _ { t } + W _ { h o } ^ { } \mathrm { H } _ { t - 1 } + W _ { c o } \circ \mathrm { C } _ { t } + b _ { o } \right)\\ \mathrm { H } _ { t } &= o _ { t } \circ \tanh \left( \mathrm { C } _ { t } \right) \end{aligned} \end{equation} $$

ConvLSTM 的公式:

$$ \begin{equation} \begin{aligned} i _ { t } &= \sigma \left( W _ { x i } * \mathrm { X } _ { t } + W _ { h i } * \mathrm { H } _ { t - 1 } + W _ { c i } \circ \mathrm { C } _ { t - 1 } + b _ { i } \right)\\ f _ { t } &= \sigma \left( W _ { x f } * \mathrm { X } _ { t } + W _ { h f } * \mathrm { H } _ { t - 1 } + W _ { c f } \circ \mathrm { C } _ { i - 1 } + b _ { f } \right)\\ C _ { t } &= f _ { t } \circ C _ { t - 1 } + i _ { t } \circ \tanh \left( W _ { x c } * X _ { t } + W _ { h c } * H _ { t - 1 } + b _ { c } \right)\\ o _ { t } &= \sigma \left( W _ { x o } * \mathrm { X } _ { t } + W _ { h o } * \mathrm { H } _ { t - 1 } + W _ { c o } \circ \mathrm { C } _ { t } + b _ { o } \right)\\ \mathrm { H } _ { t } &= o _ { t } \circ \tanh \left( \mathrm { C } _ { t } \right) \end{aligned} \end{equation} $$

$*$ 是卷积操作,$\circ$ 是哈达玛乘积。因此 FC-LSTM 与 ConvLSTM 的主要差别在于将矩阵乘法替换成了卷积操作,即 input-to-state 和 state-to-state 的转换中使用的是卷积而不是全连接。

If we view the states as the hidden representations of moving objects, a ConvLSTM with a larger transitional kernel should be able to capture faster motions while one with a smaller kernel can capture slower motions. Also, if we adopt a similar view as [16], the inputs, cell outputs and hidden states of the traditional FC-LSTM represented by (2) may also be seen as 3D tensors with the last two dimensions being 1. In this sense, FC-LSTM is actually a special case of ConvLSTM with all features standing on a single cell.

通过使用一个大的状态转换卷积核捕捉快速的运动,而小的卷积核捕捉慢的运动;当最后两个维度为 1 时,FC-LSTM 也可以视为一个特殊的 ConvLSTM,每一个单元格代表一个特征。

Pytorch 实现可以参考 ConvLSTM,我合并了八次卷积操作,所以整个计算量和普通的卷积差不多,偏置项影响不大。

TrajGRU

TrajGRU(Trajectory GRU) 来自于施行健的第二篇文章 Deep Learning for Precipitation Nowcasting: A Benchmark and A New Model ,在这篇文章中,作者提出了一个新的 Benchmark 也就是上面介绍的数据集,对比了 ConvGRU、不使用/使用加权损失函数训练、Conv2d、Conv3d 等几种方法,构建了一个 Encoder-Forecaster 模型。

ConvGRU 类似于上面介绍的 ConvLSTM,只是将 LSTM 替换成了 GRU。

而 TrajGRU 使用 GRU 结构,区别是首先使用了一个 $32@5k2p1s$ 的卷积层对隐藏态进行卷积提取光流,然后使用 $2L@1k1s$ 的卷积层对光流进行卷积,建立一个动态的对应关系。

按照作者的说法,普通卷积网络的对应关系不变,对于 local-invariant 的操作比较合适,而 local-variant 的操作比如旋转普通卷积是不够的。相比普通卷积固定的领域集,TrajGRU 可以有 $L$ 个变化的领域集,并且参数要少于原来的普通卷积。我的想法是云层的运动与消散是动态的,一团云在某一帧的位置对应到下一帧的位置不一样,所以需要建立一个动态的链接。

TrajGRU 的公式如下,首先使用一个子网络提取光流 $\mathcal{U}_t,\mathcal{V}_t$,然后使用 wrap 函数变换,有 $L$ 个链接。

$$ \begin{align} {{\mathcal{U}}_{t}},{{\mathcal{V}}_{t}} &=\gamma \left( {{\mathcal{X}}_{t}},{{\mathcal{H}}_{t-1}} \right) \\ {{\mathcal{Z}}_{t}} &=\sigma \left( {{\mathcal{W}}_{xz}}*{{\mathcal{X}}_{t}}+\sum\limits_{l=1}^{L}{\mathcal{W}_{hz}^{l}}* warp \left( {{\mathcal{H}}_{t-1}},{{\mathcal{U}}_{t,l}},{{\mathcal{V}}_{t,l}} \right) \right) \\ {{\mathcal{R}}_{t}} &=\sigma \left( {{\mathcal{W}}_{xr}}*{{\mathcal{X}}_{t}}+\sum\limits_{l=1}^{L}{\mathcal{W}_{hr}^{l}}* warp \left( {{\mathcal{H}}_{t-1}},{{\mathcal{U}}_{t,l}},{{\mathcal{V}}_{t,l}} \right) \right) \\ \mathcal{H}_{t}^{\prime } &=f\left( {{\mathcal{W}}_{xh}}*{{\mathcal{X}}_{t}}+{{\mathcal{R}}_{t}}{}^\circ \left( \sum\limits_{l=1}^{L}{\mathcal{W}_{hh}^{l}}* warp \left( {{\mathcal{H}}_{t-1}},{{\mathcal{U}}_{t,l}},{{\mathcal{V}}_{t,l}} \right) \right) \right) \\ {{\mathcal{H}}_{t}} &=\left( 1-{{\mathcal{Z}}_{t}} \right){}^\circ \mathcal{H}_{t}^{\prime }+{{\mathcal{Z}}_{t}}{}^\circ {{\mathcal{H}}_{t-1}} \end{align} $$

wrap 函数的 mxnet 可以参考 traj_rnn.py#L14,pytorch 没有该函数,可以参考我的实现 trajGRU.py#L9

Encoder-Forecaster

这两篇文章都采用了一个 Encoder-Forecaster 模型,和普通的 Encoder-Decoder 模型还是不一样的。Encoder-Forecaster 分成两部分,分别是 Encoder(编码器),Forecaster(预测器),分别对应了 Encoder 和 Decoder。

模型展开之后示意图如下:

这是一个输入 2 帧,预测未来 2 帧的 RNN 展开后的结构图。Encoder-Forecaster 模型对应层次的 RNN 的 隐藏态 输入到下一层次之中,提取不同层次的时空信息。Downsample 和 Unsample 分别通过 Convolution 和 Deconvolution 实现。这里的 RNN 可以采用 ConvGRU、ConvLSTM 和 TrajGRU 等能提取时空信息的结构。

Encoder 的初始隐藏态和 Forecaster 的初始输入为 0,因为这两个没有输入。最后的输出则通过一个 $1\times 1$ 的卷积层进行回归。

Deep Learning for Precipitation Nowcasting: A Benchmark and A New Model 的网络结构可以参考下图:

通过实验发现,TrajGRU 取得了比 ConvGRU、Conv3d、光流法等方法更好的结果;所有深度学习方法都取得了比其他方法更好地结构;文中的 Baseline 是直接使用最后一帧的作为未来所有帧的预测结果。

在训练时作者采用了一个 50 的梯度阶段,防止梯度爆炸。 还有一些更具体的参数,比如网络结构,学习率等可以参考我的实现 experiments/net_params.py

总结

我实现了 ConvLTSM 和 TrajGRU,取得了和作者接近的结果,有个疑问则是 ConvLTSM 取得了要比 TrajGRU 的更好地结果,而且因为合并了八个卷积所以运行速度也快得多。更具体的说明可以参考 关于 TrajGRU 的意义

运行效果图如下:

虽然结果可能不错,在实际预测时还是发现了越往后的输出越暗,即越往后误差越大,参考 回复

如果更进一步改进,可以考虑对这一方面、运行速度的比较等进行论述。现在气象台采用的还是光流法,预测之后由专业人员进行判断是否有下雨天气。