一种深度学习系统可以从高精度的气候模拟中识别出极端天气模式

近年来极端天气所造成的影响越来越大，不久前的山竹和飞燕带来的损失让我们对自然的力量仍然记忆犹新。如何更好的识别极端天气，做好防灾减灾工作是气象人员面临的巨大挑战。

而最近，为了帮助识别极端天气的潜在模式，预测出有可能损害人民群众生命财产安全的天气模式，来自伯克利大学劳伦斯实验室、橡树岭实验室和英伟达的研究人员们，利用AI和超级计算机给出了答案。他们研发了一种深度学习系统可以从高精度的气候模拟中识别出极端天气模式，这种算法有望在未来帮助公众尽早预知天气，在天灾面前从容应对。

研究人员们利用Tiramisu和DeepLabv3神经网络为极端天气构建了像素级的掩膜。在最近放出的论文里，研究人员详细描述了在软件架构上的改进，包括输入流程和训练算法有效的提升了在超算上的大规模深度学习计算能力。Tiramisu网络被大规模的部署在5300个P100的英伟达GPU上，实现了21PF/s的运算能力以及79%的并行效率。同时DeepLabv3则部署到了27360个V100GPU上，实现了325.8PF/s的速度，同时在单精度下实现了90.7%的并行效率。同时通过高效的使用FP16张量核心，半精度版本的DeepLabv3+网络可以实现1.13EF/s的峰值运算和999.0 PF/s的持续运算能力。这些计算在目前排名第一的超算Summit上进行。

研究人员表示这一研究实现了很多记录，不仅是在气象科学领域第一次将深度学习用于解决如此大规模的气候分割问题，同时在深度学习领域也是第一次将计算能力拓展到了Exa（艾10^18）量级。1.13EF/s的计算能力是深度学习首次跨越Exa的大门。值得一提的是，这一突破也获得了今年的Gordon Bell奖，用于鼓励研究团队在超算领域的杰出成就。

下面让我们来看看研究人员使用的网络架构，用于对天气状况进行像素级的分割。

这项工作使用改进DeepLabv3+网络，用于对气候数据进行分割。其中编码器部分使用了ResNet-50的核心部分来提取输入数据的高维特征，并使用空洞空间金字塔池化来处理较大的输入分辨率。而其解码器则被完整分辨率的统一操作取代，以便生成精确的边界信息。在上图中，深蓝色框代表了标准的卷积操作，浅蓝色代表了反卷积操作，而空洞卷积操作则由绿色代表。

研究人员们使用了CommunityAtmosphere Model (CAM5) 来生成100年的模拟数据，并用启发式的算法来生成掩膜标签。整个数据集包含了63k的高分辨率样本，其中包含了水蒸气、风力、降水、温度和气压等，分割区域包含了热带气旋、大气长河和背景三种区域。

为了在超算上大规模的部署应用，研究人员们测试了不同数量GPU的计算速度：

结果显示目前的算法和虚线表示的结果十分接近。

下图是最终对于全球范围内的预测掩膜结果，我们可以看到预测出了结果和黑色的标记结果十分接近。其中蓝色为大气长河红色为热带气旋。

这一工作在先前气象模式分割的启发下成功地将Tiramisu和DeepLabV3+网络应用到了高分辨率多变量的气候数据中，并在对损失函数、优化机制和网络结构的改进下，在定量和定性方面都实现了优异的结果。同时，为了大规模部署算法研究人员们建立了系统级的优化方式，包括了先进的数据存储策略、优化吞吐和层级通讯，使得峰值节点4560情况下利用27360个VoltaGPU实现了高效的并行计算和超高的算力（峰值1.13EF/s，均值999.0PF/s），这一研究拓展了开源工具Tensorflow和Horovod的使用，研发的系统已经部署在了Summit超算中。这将对深度学习领域带来诸多贡献，同时也将促进超算平台的发展。