模型实践 | AAAI 2022 气象预测之时空图神经网络

实验｜Vachel       算力支持｜幻方AIHPC

气象预测直接关乎人类社会的生产生活，一直以来是人类重点研究的科学领域之一。传统的天气预报采用非常复杂的数值气象模型，近年来越来越多的机器学习，深度学习方法被应用到这一领域，他们将气象预测抽象成一个时空预测问题，尝试通过多维时间序列、图神经网络（GNN）等方法来解决。

AAAI 2022前不久公布了今年的 Accept Paper List，其中来自西湖大学的几位研究者联合发表了一项科研成果。他们针对近年来交通流预测任务下的图神经网络模型，认为其架构不能直接照搬进气象预测领域并进行了证明。提出了一系列改进方法，最终模型取得了SOTA的效果。

本期模型试跑将为大家体验该项研究成果，来看看最近火热的GNN是否可以有效提升气象预测的精度，缩短预测时间。

High-flyer AIHPC 萤火二号使用申请

论文标题：Conditional Local Convolution for Spatio-temporal Meteorological Forecasting

论文地址：https://arxiv.org/abs/2101.01000

论文源码：https://github.com/BIRD-TAO/CLCRN

01 模型介绍

在交通流预测领域中，“GNN+RNN”的方法近年来获得了许多成功，其静态切片空间的图结构抽象，到动态时间上的embedding累积，有效解决了时空预测场景下的图动态问题。与交通流预测类似，气象预测也可以抽象为一个时空预测问题。然而，只是简单的套用交通流预测的方法，往往不能取得很好的效果。

气象预测领域的特殊性有：

1. 数据的不规则性：分布世界各地、太空卫星的各种气象传感器所捕捉的气象数据往往是不规则的，使得经典的CNN无法适用。

2. 非平面的数据信号：气象数据通常是立体的球形信号，而非平面信号；陆地和海洋的温度传感器各有不同，不在固定的网格结构中，但气象预测往往无法将其割裂开来单独做分析。

3. 高度的时空依赖：不同地形表现出完全不同的风流或温度转移模型，而极端气候事件往往使得数据不符合平稳性条件。

如上图所示，气候预测任务与交通流预测任务的不同在于，每一个局部点受其四面八方的因素影响，局部空间下模式应该基本类似，因为气象学中，热和风是自由扩散的。而交通流预测任务中，两个靠近的交通枢纽之间模式因为交通流向会存在很大差异。

因此，研究者们需要找到一种方法，来抽象气象预测场景下的问题。具体的：

1. 将分布地球表面的各种气象信号抽象为一个球数据结构，给予坐标，球体上要区分若干个局部空间（local space），他们与实际地理测地线、欧氏距离相关；

2. 通过单位速度测地线（unit-speed geodesics）定义出每个局部空间的邻居，进而构建起等轴映射（isometric map），将球体数据结构转化成图数据结构，以描述不同局部空间之间的关系；

针对上述问题抽象，该文研究者们提出了因地制宜的图卷积核（location-characterized kernel）来改良DCRNN（一种用于交通流预测的图神经网络，ICLR 2018）中的图卷积部分。具体的：

1. 如上图a，为了能更好的保留地球表面地理标线的相对方位，作者采用圆柱切线空间（cylindrical-tangent space）15和水平映射（horizon map）方法构建起局部空间并映射其邻居，构建起等轴映射；

2. 为了更好了模拟气象学中的局部一致性，作者重新设计了条件局部卷积核，以满足因地制宜的卷积特征，邻近局部特征相似和地理特性不同下的相邻卷积核共享三个条件。如上图b，来自东方的xp对西边的两个局部空间 i 和 j 的影响应该是一样的，不会因为xp离 j 近就在图卷积中被赋予更大的信息量；

3. 如上图c，两个相邻空间各自拥有不同的地理分布，这时不能简单的采用平滑方法以满足邻近特征相似。作者提出结合其距离和角度来重新计算卷积核的权重。具体的，在计算测地线距离的基础上，使用极坐标重新刻画数据，计算相邻点的角平分线。

最后，模型整体的结构如下图：

 02 模型实践

作者开源了实验代码，包括baseline，其主要采用图神经网络库torch-geometric进行进一步的封装实现。代码由于是刚刚开源，需要一定改造方可运行起来。目前数据和代码已整合进幻方AI的数据仓库与模型仓库当中，供后续研究者们调用。

1.数据集

该研究采用的是全球天气预报基准数据集 WeatherBench，由大数据地球研究社区 Pangeo 于2020年开源。其一共由2048个分布在地球各个重要位置上的统计节点。作者选取温度、湿度、云量、地表风分量四个天气预报项进行试验，以小时为单位划分数据，以过去12小时数据进行分析，对未来12小时进行预报。

值得注意的是，在开源的代码中，作者是将这四个预报项单独进行预测，分别测试准确度，而不是综合多个因素进行分析。例如，对温度的预测只考虑过去温度的变化，不会将过去湿度，云量，风量等因素综合考虑其中。

除了时序化的指标数据，作者还抽析出了每个指标每个时刻的地理位置信息，以构建一个球体结构。如上一章所介绍的，这个结构将被解析成一个动态图，进行时空分析并预测。所抽析的地理指标包括：经度、纬度、海拔、陆地标记。最终，组成的输入输出形式为：

2.模型训练

作者开源的代码采用的单机单卡训练模式，提供了预处理好的训练样本与测试样本，可以通过pickle直接加载使用。在幻方的AIHPC上，用户可以获得更多的高性能显卡，因此可以进行并行训练以获取更好的性能，相关内容可以参考之前的文章《PyTorch分布式训练方法》。

这里，我们采用分布式数据并行（DDP）的方式改造代码，提高训练性能。具体的：

 1. 用 DistributedDataParallel 封装原始模型

from torch.nn.parallel import DistributedDataParallel

self.model = DistributedDataParallel(clcrn_model.cuda(), device_ids=[local_rank])

2. 用 DistributedSampler 封装原始数据集，使其自由分散在多张显卡中

from torch.utils.data.distributed import DistributedSampler

train_datasampler = DistributedSampler(trainset)
self.trainloader = torch.utils.data.DataLoader(trainset, batch_size, sampler=train_datasampler, num_workers=4, pin_memory=True)

3. 加入 hfai 监听集群中的统一调度信号，做好checkpoint 保存

import hfai

# the first node gets the suspend signal, save checkpoint
if rank == 0 and hfai.receive_suspend_command():
    self.save_model("checkpoint.m")
    time.sleep(5)
    hfai.go_suspend()

4. 通过 Process 托管多进程，管理多张显卡上的训练

import torch

if __name__ == '__main__':
    torch.multiprocessing.spawn(main, args=(), nprocs=torch.cuda.device_count())

改造完成后，运行的效果如下：

主节点Epoch平均耗时48s左右，相比单机单卡训练提升速度6-7倍。完整的改造代码已集成进幻方AI的模型仓库中，感兴趣的研究者们可以滑至文末申请超算体验，了解更多的技术细节。

类似WeatherBench中展示的可视化案例，可以将预测出的天气信息映射到地图中，对比出天气的变化情况,。如下图所展示的温度变化案例：

03 体验总结

气象预测与人类生产生活息息相关。AAAI 2022这篇研究工作将近年来火热的时空图神经网络应用到了气象预测当中，将这一问题具体抽象成了图数据结构并予以解决。该工作进一步推进了AI+环境科学的发展。同时，无论是模型还是该场景中的问题，还有许多能进一步研究的空间。比如，是否可以结合更多源的数据对某一预报项进行预测、不同因素之间彼此是否有关联、扩大数据规模增加算力能提升多少精度等等。我们欢迎感兴趣的学者加入我们，继续在幻方AI团队的萤火超算平台上对该课题进行进一步探索。

综合体验打分如下：

01：研究新颖度  ★★★★

气象预测是重要的研究课题，结合AI方法有很多的创新空间。

02：开源指数  ★★★★

数据和代码都进行了完整开源，代码部分修改后可以运行。

03：算力门槛  ★★★

作者开源的数据规模不大，对算力的需求较小。但随着数据规模的提升和多指标的整合分析，算力的需求会逐渐增大。

04：通用指数  ★★

模型针对气象预测领域的特性进行模型设计，通用性比较差。

05：模型适配度  ★★

模型采用单机单卡方式实现，需要进行多机并行训练改良方可适配幻方AIHPC。

幻方AI团队， 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI实验落地与价值创造，降低超大算力需求的使用门槛，欢迎各方数据科学家与开发者们一同共建。