FourCastNet: 基于傅里叶神经算子的全球中期天气预报

概述

FourCastNet（傅里叶预测神经网络）是一个数据驱动的全球天气预报模型，由NVIDIA、劳伦斯伯克利国家实验室、密歇根大学安阿伯和赖斯大学的研究人员开发。它提供了关键全球天气指标的中期预报，分辨率为0.25°。相当于赤道附近约30公里x30公里的空间分辨率和大小为720 x 1440像素的全球网格。与传统的NWP模型相比，该模型的预报速度提高了45000倍，2秒内生成一周的天气预报，预报精度与最先进的数值天气预报模型ECMWF综合预报系统（IFS）相当。这是第一个可以直接与IFS系统进行比较的AI天气预报模型。

本教程介绍了FourCastNet的研究背景和技术路径，并展示了如何通过MindFlow训练和快速推断模型。更多详细信息见文章。

技术路径

MindEarth求解该问题的具体流程如下：

创建数据集
模型构建
损失函数
模型训练
模型验证和可视化

FourCastNet

为了实现高分辨率预测，FourCastNet使用AFNO模型。该模型网络体系结构是为高分辨率输入而设计的，以ViT为骨干网，并结合了李宗义等人提出的傅里叶神经算子(FNO)。该模型学习函数空间之间的映射，从而求解一系列非线性偏微分方程。

Vision Transforme（ViT）体系结构及其变体在过去几年中已成为计算机视觉中最先进的技术，在许多任务中表现出卓越的性能。这种性能主要归因于网络中的多头自注意机制，它使网络中每一层特征之间的全局建模。然而，模型在训练和推理期间的计算复杂度随着令牌（或patches）数量的增加而二次增加，模型计算复杂度随着输入分辨率的增加而爆炸性增加。

AFNO模型的独创性在于，它将空间混合操作转换为傅里叶变换，混合不同令牌的信息，将特征从空域转换为频域，并对频域特征应用全局可学习滤波器。空间混合复杂度有效地降低到O(NlogN)，其中N是token的数量。

FourCastNet网络架构如下图所示。

AFNO model

模型训练包括三个步骤：

1.预训练：如上图（a）所示，在预训练步骤中，使用训练数据集以监督的方式训练AFNO模型，以学习从X(k)到X(k + 1)的映射。

2.微调：如上图（b）所示，模型首先从X(k)预测X(k + 1)，然后使用X(k + 1)作为输入预测X(k + 2)。然后，通过从X(k + 1)和X(k + 2)的预测值计算损失函数值，利用两个损失函数值的总和优化模型。

3.降水预报：如上文（c）所示，降水预报由主干模型后面的单独模型拼接而成。该方法将降水预测任务与基本气象要素解耦。另一方面，训练好的降水模型也可以与其他预测模型（传统的NWP等）结合使用。

本教程主要实现模型预训练部分。

[1]:

import os
import numpy as np
import matplotlib.pyplot as plt

from mindspore import context
from mindspore import load_checkpoint, load_param_into_net

from mindearth.utils import load_yaml_config, create_logger, plt_global_field_data
from mindearth.module import Trainer
from mindearth.data import Dataset, Era5Data
from mindearth import RelativeRMSELoss
from mindearth.cell import AFNONet

src 文件可以从FourCastNet/src下载。

[2]:

from src.callback import EvaluateCallBack, InferenceModule

context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=2)

model、data和optimizer的参数可以通过加载yaml文件获取（FourCastNet.yaml）。

[3]:

config = load_yaml_config('FourCastNet.yaml')
config['model']['data_sink'] = True  # 是否使用data sink特性

config['train']['distribute'] = False  # 是否执行分布式任务
config['train']['amp_level'] = 'O2'  # 设置混合精度等级

config['data']['num_workers'] = 1  # 设置并行计算的进程数量
config['data']['grid_resolution'] = 1.4  # 设置气象分辨率参数

config['optimizer']['epochs'] = 100  # 设置epoch数量
config['optimizer']['finetune_epochs'] = 1  # 设置微调epoch数量
config['optimizer']['warmup_epochs'] = 1  # 设置预热epoch的数量
config['optimizer']['initial_lr'] = 0.0005  # 设置初始化学习率

config['summary']["valid_frequency"] = 10  # 设置验证的频率
config['summary']["summary_dir"] = './summary'  # 设置模型checkpoint的存储路径

logger = create_logger(path="results.log")

创建数据集

在dataset路径下，下载正则化参数、训练数据集验证数据集到 ./dataset目录。

修改FourCastNet.yaml配置文件中的root_dir参数，该参数设置了数据集的路径。

./dataset中的目录结构如下所示：

.
├── statistic
│   ├── mean.npy
│   ├── mean_s.npy
│   ├── std.npy
│   └── std_s.npy
├── train
│   └── 2015
├── train_static
│   └── 2015
├── train_surface
│   └── 2015
├── train_surface_static
│   └── 2015
├── valid
│   └── 2016
├── valid_static
│   └── 2016
├── valid_surface
│   └── 2016
├── valid_surface_static
│   └── 2016

模型构建

加载相关的数据参数和模型参数，并完成AFNONet模型构建。

[4]:

data_params = config['data']
model_params = config['model']

model = AFNONet(image_size=(data_params['h_size'], data_params['w_size']),
                in_channels=data_params["feature_dims"],
                out_channels=data_params["feature_dims"],
                patch_size=data_params["patch_size"],
                encoder_depths=model_params["encoder_depths"],
                encoder_embed_dim=model_params["encoder_embed_dim"],
                mlp_ratio=model_params["mlp_ratio"],
                dropout_rate=model_params["dropout_rate"])

损失函数

FourCastNet使用相对均方根误差进行模型训练。

[5]:

loss_fn = RelativeRMSELoss()

模型训练

在本教程中，我们继承了Trainer并重写了get_callback成员函数，以便我们可以在训练过程中对测试数据集执行推理。

对于MindSpore版本>= 1.8.1，我们可以使用函数式编程来训练神经网络。MindSpore Earth为模型训练提供了训练接口。

[6]:

class FCNTrainer(Trainer):
    def __init__(self, config, model, loss_fn, logger):
        super(FCNTrainer, self).__init__(config, model, loss_fn, logger)
        self.pred_cb = self.get_callback()

    def get_callback(self):
        pred_cb = EvaluateCallBack(self.model, self.valid_dataset, self.config, self.logger)
        return pred_cb

trainer = FCNTrainer(config, model, loss_fn, logger)

2023-09-07 02:26:03,143 - pretrain.py[line:211] - INFO: steps_per_epoch: 404

[7]:

trainer.train()

epoch: 1 step: 404, loss is 0.5348429
Train epoch time: 136480.515 ms, per step time: 337.823 ms
epoch: 2 step: 404, loss is 0.35937342
Train epoch time: 60902.627 ms, per step time: 150.749 ms
...
epoch: 98 step: 404, loss is 0.15447393
Train epoch time: 61055.706 ms, per step time: 151.128 ms
epoch: 99 step: 404, loss is 0.15696357
Train epoch time: 60850.156 ms, per step time: 150.619 ms
epoch: 100 step: 404, loss is 0.15654306
Train epoch time: 60944.369 ms, per step time: 150.852 ms
2023-09-07 04:27:02,837 - forecast.py[line:209] - INFO: ================================Start Evaluation================================
2023-09-07 04:28:25,277 - forecast.py[line:177] - INFO: t = 6 hour:
2023-09-07 04:28:25,277 - forecast.py[line:188] - INFO:  RMSE of Z500: 154.07894852240838, T2m: 2.0995438696856965, T850: 1.3081689948838815, U10: 1.527248748050362
2023-09-07 04:28:25,278 - forecast.py[line:189] - INFO:  ACC  of Z500: 0.9989880649296732, T2m: 0.9930711917863625, T850: 0.9954355203713009, U10: 0.9615764420500764
2023-09-07 04:28:25,279 - forecast.py[line:177] - INFO: t = 72 hour:
2023-09-07 04:28:25,279 - forecast.py[line:188] - INFO:  RMSE of Z500: 885.3778200063341, T2m: 4.586325958437852, T850: 4.2593739999338736, U10: 4.75655467109408
2023-09-07 04:28:25,280 - forecast.py[line:189] - INFO:  ACC  of Z500: 0.9598951919101183, T2m: 0.9658168304842388, T850: 0.9501612262744354, U10: 0.6175327930007481
2023-09-07 04:28:25,281 - forecast.py[line:177] - INFO: t = 120 hour:
2023-09-07 04:28:25,281 - forecast.py[line:188] - INFO:  RMSE of Z500: 1291.3199606908572, T2m: 6.734047767054735, T850: 5.6420206614200294, U10: 5.637643311177468
2023-09-07 04:28:25,282 - forecast.py[line:189] - INFO:  ACC  of Z500: 0.9150022892106006, T2m: 0.9294266102808937, T850: 0.9148957221265037, U10: 0.47971871343985495
2023-09-07 04:28:25,283 - forecast.py[line:237] - INFO: ================================End Evaluation================================

模型推理及可视化

完成训练后，我们使用第100个ckpt进行推理。

[8]:

pred_time_index = 0

params = load_checkpoint('./summary/ckpt/step_1/FourCastNet_1-100_404.ckpt')
load_param_into_net(model, params)
inference_module = InferenceModule(model, config, logger)

[9]:

def plt_data(pred, label, root_dir, index=0):
    """ Visualize the forecast results """
    std = np.load(os.path.join(root_dir, 'statistic/std.npy'))
    mean = np.load(os.path.join(root_dir, 'statistic/mean.npy'))
    std_s = np.load(os.path.join(root_dir, 'statistic/std_s.npy'))
    mean_s = np.load(os.path.join(root_dir, 'statistic/mean_s.npy'))
    pred, label = pred[index].asnumpy(), label.asnumpy()[..., index, :, :]
    plt.figure(num='e_imshow', figsize=(100, 50), dpi=100)

    plt.subplot(4, 3, 1)
    plt_global_field_data(label, 'Z500', std, mean, 'Ground Truth')  # Z500
    plt.subplot(4, 3, 2)
    plt_global_field_data(pred, 'Z500', std, mean, 'Pred')  # Z500
    plt.subplot(4, 3, 3)
    plt_global_field_data(label - pred, 'Z500', std, mean, 'Error')  # Z500

    plt.subplot(4, 3, 4)
    plt_global_field_data(label, 'T850', std, mean, 'Ground Truth')  # T850
    plt.subplot(4, 3, 5)
    plt_global_field_data(pred, 'T850', std, mean, 'Pred')  # T850
    plt.subplot(4, 3, 6)
    plt_global_field_data(label - pred, 'T850', std, mean, 'Error')  # T850

    plt.subplot(4, 3, 7)
    plt_global_field_data(label, 'U10', std_s, mean_s, 'Ground Truth', is_surface=True)  # U10
    plt.subplot(4, 3, 8)
    plt_global_field_data(pred, 'U10', std_s, mean_s, 'Pred', is_surface=True)  # U10
    plt.subplot(4, 3, 9)
    plt_global_field_data(label - pred, 'U10', std_s, mean_s, 'Error', is_surface=True)  # U10

    plt.subplot(4, 3, 10)
    plt_global_field_data(label, 'T2M', std_s, mean_s, 'Ground Truth', is_surface=True)  # T2M
    plt.subplot(4, 3, 11)
    plt_global_field_data(pred, 'T2M', std_s, mean_s, 'Pred', is_surface=True)  # T2M
    plt.subplot(4, 3, 12)
    plt_global_field_data(label - pred, 'T2M', std_s, mean_s, 'Error', is_surface=True)  # T2M

    plt.savefig(f'pred_result.png', bbox_inches='tight')
    plt.show()

[12]:

test_dataset_generator = Era5Data(data_params=config["data"], run_mode='test')
test_dataset = Dataset(test_dataset_generator, distribute=False,
                       num_workers=config["data"]['num_workers'], shuffle=False)
test_dataset = test_dataset.create_dataset(config["data"]['batch_size'])
data = next(test_dataset.create_dict_iterator())
inputs = data['inputs']
labels = data['labels']
pred = inference_module.forecast(inputs)
plt_data(pred, labels, config['data']['root_dir'])

下述展示了第100个ckpt的真实值、预测值和他们之间的误差可视化。

plot result