当前位置：首页 > news >正文

Aurora模型数据准备指南：如何正确构建Batch对象进行预测

news 2026/6/13 6:34:46

Aurora模型数据准备指南：如何正确构建Batch对象进行预测

【免费下载链接】auroraImplementation of the Aurora model for Earth system forecasting项目地址: https://gitcode.com/gh_mirrors/aurora25/aurora

Aurora模型是一款强大的地球系统预测工具，而构建正确的Batch对象是确保预测准确性的关键步骤。本指南将带您全面了解Batch对象的核心概念、构建方法和最佳实践，帮助您快速掌握数据准备的精髓。

什么是Batch对象？

Batch对象是Aurora模型处理数据的基本单元，它封装了进行地球系统预测所需的各类数据和元信息。在Aurora模型的架构中，Batch对象扮演着数据容器的角色，负责在模型各个组件间传递和处理数据。

从源码实现来看，Batch类定义在aurora/batch.py文件中，主要包含以下四个核心部分：

surf_vars：地表变量字典，形状为(b, t, h, w)
static_vars：静态变量字典，形状为(h, w)
atmos_vars：大气变量字典，形状为(b, t, c, h, w)
metadata：元数据信息，包含经纬度、时间、气压层等关键信息

图1：Aurora模型数据结构示意图，展示了地球系统预测中各类变量的组织方式

Batch对象的核心组件解析

Metadata元数据

Metadata是Batch对象的重要组成部分，它包含了数据的关键上下文信息。根据aurora/batch.py的定义，Metadata主要包括：

lat/lon：经纬度数据，必须满足纬度严格递减、经度严格递增的要求
time：时间元组，每个批次元素的时间信息
atmos_levels：气压层信息，单位为hPa
rollout_step：滚动预测步数，默认为0表示实际数据

元数据的正确设置对模型预测的准确性至关重要，尤其是经纬度的有效性验证会在__post_init__方法中自动进行。

变量类型详解

Batch对象包含三种主要变量类型：

地表变量(surf_vars)：如温度、降水等地表观测数据
静态变量(static_vars)：如地形高度等不随时间变化的变量
大气变量(atmos_vars)：不同气压层的大气数据，具有额外的高度维度

这些变量在模型中会经过aurora/normalisation.py中的标准化处理，以确保模型能够稳定训练和预测。

构建Batch对象的完整步骤

1. 准备原始数据

首先需要收集和整理地球系统预测所需的各类数据，包括地表变量、静态变量和大气变量。这些数据通常以NetCDF格式存储，可以从气象数据中心或相关研究机构获取。

2. 创建Metadata对象

根据数据的实际情况创建Metadata对象，确保经纬度范围和气压层信息准确无误：

from aurora.batch import Metadata import torch lat = torch.tensor([90.0, 89.0, ..., -90.0]) # 严格递减的纬度 lon = torch.tensor([0.0, 1.0, ..., 359.0]) # 严格递增的经度 time = (datetime(2023, 1, 1), datetime(2023, 1, 2)) # 时间元组 atmos_levels = (1000, 850, 700, 500, 300, 200, 100) # 气压层 metadata = Metadata(lat=lat, lon=lon, time=time, atmos_levels=atmos_levels)

3. 构建变量字典

将准备好的数据组织成符合要求的变量字典：

# 地表变量字典 surf_vars = { "temperature": torch.randn(1, 3, 181, 360), # (批次, 时间, 高度, 宽度) "precipitation": torch.randn(1, 3, 181, 360) } # 静态变量字典 static_vars = { "elevation": torch.randn(181, 360) # (高度, 宽度) } # 大气变量字典 atmos_vars = { "temperature": torch.randn(1, 3, 7, 181, 360), # (批次, 时间, 气压层, 高度, 宽度) "wind_u": torch.randn(1, 3, 7, 181, 360), "wind_v": torch.randn(1, 3, 7, 181, 360) }

4. 实例化Batch对象

最后，使用上述组件创建完整的Batch对象：

from aurora.batch import Batch batch = Batch( surf_vars=surf_vars, static_vars=static_vars, atmos_vars=atmos_vars, metadata=metadata )

从文件加载Batch对象的快捷方法

Aurora模型提供了直接从NetCDF文件加载Batch对象的便捷方法，无需手动构建各个组件：

# 从NetCDF文件加载Batch batch = Batch.from_netcdf("path/to/your/data.nc") # 将Batch保存为NetCDF文件 batch.to_netcdf("path/to/save/data.nc")

这个功能在aurora/batch.py中实现，通过xarray库处理NetCDF文件的读写操作。使用这种方法可以大大简化数据准备流程，特别适合处理大型地球系统数据集。

Batch对象的关键操作与转换

数据标准化与反标准化

Batch对象提供了内置的标准化和反标准化方法，确保数据符合模型的输入要求：

# 标准化Batch数据 normalised_batch = batch.normalise(surf_stats={"temperature": (280.0, 30.0)}) # 反标准化Batch数据 unnormalised_batch = normalised_batch.unnormalise(surf_stats={"temperature": (280.0, 30.0)})

标准化过程会对地表变量、静态变量和大气变量分别应用不同的标准化策略，具体实现可参考aurora/normalisation.py。

数据裁剪与重采样

根据模型要求调整数据空间分辨率：

# 裁剪数据以适应模型输入大小 cropped_batch = batch.crop(patch_size=128) # 重采样数据到指定分辨率 resampled_batch = batch.regrid(res=0.5) # 0.5度分辨率

这些操作对于不同来源的数据整合非常有用，能够确保输入模型的数据具有一致的空间尺度。

设备迁移与类型转换

在模型训练和推理过程中，可能需要在不同设备间移动数据或转换数据类型：

# 将Batch移动到GPU gpu_batch = batch.to("cuda:0") # 转换数据类型 float32_batch = batch.type(torch.float32)

Batch对象在预测中的应用

构建好的Batch对象是模型预测的直接输入。在Aurora模型中，Batch对象会通过aurora/rollout.py中的rollout函数进行多步预测：

from aurora import AuroraPretrained, rollout # 加载预训练模型 model = AuroraPretrained.from_pretrained("path/to/pretrained/model") # 执行多步预测 predictions = rollout(model, batch, steps=10) # 处理预测结果 for step, pred_batch in enumerate(predictions): print(f"Step {step}: Prediction for {pred_batch.metadata.time[0]}") # 可视化或保存预测结果

图2：Aurora模型台风轨迹预测示例，展示了模型对台风路径的多步预测结果