直播预告 | 1分钟预测10天全球天气！MindSpore Earth 带你玩转登顶Science的天气预报大模型

2023/12/01

直播预告 | 1分钟预测10天全球天气！MindSpore Earth 带你玩转登顶Science的天气预报大模型

活动预告：12月4日晚19:00-20:00，昇思MindSpore AI4S工程师带你看天气预报大模型，全面讲解GraphCast并介绍昇思MindSpore Earth地球科学套件的实现，欢迎各位同学参加！

GraphCast是DeepMind提出的基于图神经网络的AI中期天气预报模型，提供了全球关键气象要素的高分辨率中期预报，预报精度超越目前最先进的数值预报模式。相关研究论文以“Learning skillful medium-range global weather forecasting”为题，已发表在权威科学期刊 Science 上。昇思MindSpore Earth完整实现了GraphCast高分辨率训练与推理复现，并已开源。

背景介绍

天气现象与人类的生产生活、社会经济等方方面面都密切相关，准确的天气预报能够在灾害天气事件中减轻影响、甚至避免经济损失，还能创造持续不断地财政收入，例如能源、农业、交通和娱乐行业。从小的方面来说，每天人们在计划自己的行程时都会考虑接下来的天气，当这些决定涉及预计未来约三至十天的天气时，人们依靠“中期”天气预报进行判断。在中期天气预报领域，欧洲的ECMWF的综合预报系统（IFS）是目前最先进的数值预报模式，但数值模式需要大量超级计算机的投入，是一个耗时且昂贵的过程。而采用图神经网络（GNN）机器学习架构的GraphCast模型，仅在一分钟内便可生成10天的预报，其运行成本可节省1000倍。

GraphCast

GraphCast使用GNN作为骨干网络，基于“Encode-Process-Decode”的网络结构专为复杂的物理动力学而设计，使用消息传递机制，允许在任何范围内进行任意空间交互模式。它提供了关键全球天气指标的中期预报，分辨率为0.25°，相当于赤道附近约25kmx25km的空间分辨率。在每一个网格点上，该模型都考虑了地表的五个变量，如温度、压力、湿度、经向风速和纬向风速，以及37个不同高度的大气中的六个变量。

GraphCast模型包括编码器层、处理器层和解码器层。编码器将历史时刻气象要素的纬度-经度输入网格映射到多尺度内部网格表示；处理器在多网格表示上执行多轮消息传递；解码器将多网格表示映射回纬度-经度网格，同时输出预测结果。

Figure 1 GraphCast模型架构

GraphCast模型通过当前时间状态和上一个时间状态预测下一步天气状态。形式如下：

针对多步预测精度衰减，采用rollout多步迭代训练，降低模型误差累积，形式如下：

MindSpore Earth实践

1.数据集

ERA5是ECMWF（欧洲中期天气预报中心）发布的全球气候第五代大气再分析数据集，提供了大量大气、陆地和海洋气候变量的每小时估计值。我们在ERA5的一个子集上训练我们的模型，输入分辨率为0.25°。我们选取了13个气压层（即50hPa、100hPa、150hPa、200hPa、250hPa、300hPa、400hPa、500hPa、600hPa、700hPa、850hPa、925hPa和1000hPa）及地表变量作为输入特征。同时，MindSpore Earth开源了正二十面体网格生成模块，用户可根据需要生成不同尺度和分辨率的多层级网格。多层级网格是一个空间均质的图，最高分辨率可通过 6 次迭代正二十面体（包含12个节点，20个面和30条边）形成，每次迭代会对网格进行精细化处理，将单个三角形划分为 4 个较小的三角形，并将其节点投影至球体上。

为便于用户快速实践MindSpore Earth开源的GraphCast案例，我们提供了部分ERA5 1.4°分辨率的数据，用户可从下面链接中下载。

https://download.mindspore.cn/mindscience/mindearth/dataset/WeatherBench\_1.4\_69/

2.模型实现

MindSpore Earth提供 GraphCastNet网络接口，便于快速创建模型。

model = GraphCastNet(vg_in_channels=data_params.get('feature_dims') * data_params.get('t_in'),
                   vg_out_channels=data_params.get('feature_dims'),
                   vm_in_channels=model_params.get('vm_in_channels'),
                   em_in_channels=model_params.get('em_in_channels'),
                   eg2m_in_channels=model_params.get('eg2m_in_channels'),
                   em2g_in_channels=model_params.get('em2g_in_channels'),
                   latent_dims=model_params.get('latent_dims'),
                   processing_steps=model_params.get('processing_steps'),
                …
                recompute=model_params.get('recompute', False))

在编码器层中，所有输入特征都使用多层感知器（MLP）嵌入到隐层空间中，将原始经纬度网格映射到多层级网格。

class Encoder(nn.Cell):
    def __init__(self, vg_in_channels, vm_in_channels, em_in_channels, eg2m_in_channels…):
        super(Encoder, self).__init__()
        self.feature_embedder = Embedder(vg_in_channels, …)
       self.g2m_gnn = G2MGnn(node_in_channels=latent_dims, …)

    def construct(self, grid_node_feats, …):
        vg, vm, em, eg2m, em2g = self.feature_embedder(grid_node_feats, …)
        eg2m, vm, vg = self.g2m_gnn(eg2m, vm, vg)
        return vg, vm, em, eg2m, em2g

在处理器中，一个16层深度GNN，在多网格上执行可学习的信息传递，通过长程边使信息在空间有效传播。

class Processor (nn.Cell):
    def __init__(self, node_in_channels, node_out_channels, edge_in_channels, …):
        super(Processor, self).__init__()
        self.processing_steps = processing_steps
        self.cell_list = nn.SequentialCell()
        for _ in range(self.processing_steps):
            self.cell_list.append(InteractionLayer(node_in_channels, …))

    def construct(self, node_feats, edge_feats):
        node_feats, edge_feats = self.cell_list((node_feats, edge_feats))
        return node_feats, edge_feats

在解码器中，多层级网格信息被映射回原始经纬度网格，并输出预测结果。

class Decoder (nn.Cell):
    def __init__(self, node_in_channels, node_out_channels, edge_in_channels, …):
        super(Decoder, self).__init__()
        self.m2g_gnn = M2GGnn(node_in_channels, …)
        self.node_fn = MLPNet(in_channels=node_in_channels, …)

    def construct(self, m2g_edge_feats, mesh_node_feats, grid_node_feats):
        m2g_edge_feats, mesh_node_feats, grid_node_feats = self.m2g_gnn(m2g_edge_feats, …)
        return self.node_fn(grid_node_feats)

3.损失函数

通过自定义均方差计算模型训练的loss：

class CustomWithLossCell(nn.Cell):
    def __init__(self, backbone, loss_fn, data_params):
        super(CustomWithLossCell, self).__init__(auto_prefix=False)
        self._backbone = backbone
        …

    def construct(self, data, labels):
        pred_list = _forecast_multi_step(data, self._backbone, self.feature_dims, …)
        loss = 0
        for t in range(self.t_out_train):
            pred = pred_list[t]
            if self.t_out_train == 1:
                label = ops.squeeze(labels)
            else:
                label = ops.squeeze(labels[:, t])
            loss_step = self._loss_fn(label, pred)
            loss += loss_step
        loss = loss / self.t_out_train
        return loss

4.模型训练

基于昇思MindSpore提供的Model接口，即可快速开展模型的训练。同时，可通过重写get_callback函数，在训练过程中对模型进行推理评价。

class GraphCastTrainer(Trainer):
    def __init__(self, config, model, loss_fn, logger):
        super().__init__(config, model, loss_fn, logger)
        self.train_dataset, self.valid_dataset = self.get_dataset()
        self.pred_cb = self.get_callback()
        self.solver = self.get_solver()

    def get_solver(self):
        loss_scale = DynamicLossScaleManager()
        solver = Model(network=self.loss_fn,
                     optimizer=self.optimizer,
                     loss_scale_manager=loss_scale,
                     amp_level=self.train_params.get('amp_level', 'O2'))
        return solver

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

若采用高分辨率网格（0.25°）、精细化多层级网格(M6：包含 40,962个节点， 81920个面)， GraphCast在训练中将需要大量内存。MindSpore Earth使用重计算、混合精度、异构存储等方式实现内存优化，从而在昇腾上实现上述输入的训练：

重计算：MindSpore在计算时，可以不保存正向算子的计算结果，让这些内存可以被复用，然后在计算反向算子时，如果需要正向的结果，再重新计算正向算子，从而降低峰值。
混合精度：在训练时，对神经网络不同的运算采用不同的数值精度的运算策略。在神经网络运算中，部分运算对数值精度不敏感，此时使用较低精度其所占用内存也会相应省，还可以达到明显的加速效果。
异构存储：在训练过程中，将暂时不需要参与计算的参数或中间结果拷贝至Host侧内存，甚至硬盘存储，在需要数据参与计算时，再将其拷贝恢复至设备侧。可以提升相同硬件设备能够训练的模型规模。

5.模型评估

我们在不同规模数据集（4年、16年）上进行训练，将测试集上的预报结果同欧洲中期天气预报中心（ECMWF）的高分辨率综合预测系统（IFS）模型进行对比。对于关键气象指标Z500, T850, T2m和U10，MindSpore Earth在16年数据集上训练的GraphCast模型，预报精度已超过IFS。