速度提升45000倍,英伟达用傅里叶模型实现前所未有天气预报准确率( 二 )


在实现技术上 , FourCastNet 使用基于傅里叶变换的token混合方法 [Guibas et al., 2022] 和ViT骨干 [Dosovitskiy et al., 2021] 。 这一方法基于最近的的傅里叶神经算子 , 该算子以分辨率不变的方式学习 , 并在建模流体动力学等具有挑战性的偏微分方程中取得了成功 。 此外 , 他们选择 ViT 骨干的原因是它能够很好地建模长程依赖 。 ViT 和基于傅里叶的 token 方法混合生成了 SOTA 高分辨率模型 , 它可以解析细粒度的特征 , 并能够很好地随分辨率和数据集大小扩展 。 研究者表示 , 这一方法能够以真正前所未有的高分辨率训练高保真数据驱动的模型 。
训练模型
欧洲中期天气预报中心(ECMWF)提供了一个公开可用的综合数据集 ERA5 , 该研究使用 ERA5 来训练 FourCastNet 。 他们专注于两个大气变量 , 即(1)距离地球表面 10m 处的风速和(2)6 小时总降水量 , 除此以外 , 该研究还预测了其他几个变量 , 包括几个不同垂直高度的位势高度、温度、风速和相对湿度 , 一些近地表变量 , 如地面气压和平均海平面气压以等 。
FourCastNet模型
为了生成高分辨率预测 , 研究采用自适应傅里叶神经算子 (Adaptive Fourier Neural Operator , AFNO) 模型 。 这种神经网络架构专为高分辨率输入而设计 , 并将深度学习(DL)最新关键进展综合到一个模型中 。 也就是说 , FourCastNet 结合了 傅里叶神经算子 (FNO), 该方法已被证明在建模具有挑战性的 PDE 系统方面表现良好 , 并且还具有强大的 ViT 主干 。
目前 , 研究者已经提出了几种降低计算复杂度的 ViT 变体 , 然而 , AFNO 模型的独特之处在于它将混合操作构建为连续的全局卷积 , 在傅里叶域中通过 FFT 有效实现 , 这允许灵活且可扩展地建模跨空间和通道维度的依赖关系 。 通过这样的设计 , 空间混合复杂度降低到 O(N log N) , 其中 N 是图像 patch 或 token 的数量 。 这种缩放使 AFNO 模型非常适合 0.25° 分辨率的高分辨率数据 , 以及进行未来更高分辨率的研究 。
模型的计算流程:首先将 720 × 1440 lat-lon 网格上的输入变量投影到 patch 的 2D 网格 (h × w)(patch大小为 p × p , 例如 p = 8) , 每个 patch 表示为一个 d 维 token 。 然后 , 将 patch 序列连同位置编码一起馈送到一系列 AFNO 层 。
训练
该研究主要关注预测地表风速和降水 , 但复杂的大气系统包含多个变量之间的强非线性相互作用 , 例如温度、地面气压、湿度、从地球表面到平流层的水分含量等 。 为了模拟这些相互作用 , 该研究选择了一些变量(表 1)来表示大气的瞬时状态 。
速度提升45000倍,英伟达用傅里叶模型实现前所未有天气预报准确率
文章图片

该研究将 ERA5 数据集分为三组 , 即训练、验证和测试数据集 。 训练数据集由 1979 年至 2015 年的数据组成 , 验证数据集包含 2016 年和 2017 年的数据 , 测试数据集包含 2018 年及以后的数据 。
该研究使用张量 X(k?t) 表示变量 , 其中k为时间指数 , ?t 为训练数据集中连续快照之间的时间间隔 , 研究将 ERA5 数据集视为真值 , 用 X_true(k?t) 表示真值变量 , ?t 固定在 6 小时 。 训练包括预训练和微调两个阶段:

  • 预训练阶段 , 采用训练数据集以有监督的方式训练 AFNO 模型 , 使其学习从 X(k) 到 X(k + 1) 的映射;
  • 微调阶段 , 该研究从之前训练好的模型开始对模型进行优化 , 即模型首先从输入 X(k) 生成输出 X(k + 1) , 然后模型使用输出 X(k + 1) 作为输入 , 并生成输出 X(k + 2) 。 通过比较 X(k + 1) 和 X(k + 2) 与训练数据中各自的真值来计算训练损失 , 并使用这两个训练损失的总和来优化模型 。

    特别声明:本站内容均来自网友提供或互联网,仅供参考,请勿用于商业和其他非法用途。如果侵犯了您的权益请与我们联系,我们将在24小时内删除。