幻方

#幻方 | 来源: 网络整理| 查看: 265

超算跑模型 | AAAI 2022 气象预测之时空图神经网络

#顶尖人才招聘 #DL/RL/NLP/CV/C++/K8S/全栈......持续热招中！详情请跳转以下链接：

百万年薪！幻方开启全球招聘！

模型介绍

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

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

数据的不规则性：分布世界各地、太空卫星的各种气象传感器所捕捉的气象数据往往是不规则的，使得经典的CNN无法适用。非平面的数据信号：气象数据通常是立体的球形信号，而非平面信号；陆地和海洋的温度传感器各有不同，不在固定的网格结构中，但气象预测往往无法将其割裂开来单独做分析。高度的时空依赖：不同地形表现出完全不同的风流或温度转移模型，而极端气候事件往往使得数据不符合平稳性条件。

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

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

将分布地球表面的各种气象信号抽象为一个球数据结构，给予坐标，球体上要区分若干个局部空间（local space），他们与实际地理测地线、欧氏距离相关；通过单位速度测地线（unit-speed geodesics）定义出每个局部空间的邻居，进而构建起等轴映射（isometric map），将球体数据结构转化成图数据结构，以描述不同局部空间之间的关系；

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

如上图a，为了能更好的保留地球表面地理标线的相对方位，作者采用圆柱切线空间（cylindrical-tangent space）15和水平映射（horizon map）方法构建起局部空间并映射其邻居，构建起等轴映射；为了更好了模拟气象学中的局部一致性，作者重新设计了条件局部卷积核，以满足因地制宜的卷积特征，邻近局部特征相似和地理特性不同下的相邻卷积核共享三个条件。如上图b，来自东方的xp对西边的两个局部空间 i 和 j 的影响应该是一样的，不会因为xp离 j 近就在图卷积中被赋予更大的信息量；如上图c，两个相邻空间各自拥有不同的地理分布，这时不能简单的采用平滑方法以满足邻近特征相似。作者提出结合其距离和角度来重新计算卷积核的权重。具体的，在计算测地线距离的基础上，使用极坐标重新刻画数据，计算相邻点的角平分线。

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

模型实践

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

1. 数据集

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

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

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

2. 模型训练

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

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

1. 用 DistributedDataParallel 封装原始模型

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

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

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

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

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

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

体验总结

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

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

扫码关注「幻方AI 」，微信公众号 | HFAILAB，算力资源随时待命

【本文地址】

公司简介

联系我们