Corrformer (《自然机器智能》)

在这篇论文中，我们提出了具有多相关性机制的Corrformer，它可以在学习到的多尺度树结构中统一时间自相关性和空间相关性。

Corrformer将时空建模的标准二次复杂度降低到空间建模的线性复杂度和时间建模的对数线性复杂度，首次在一个统一的深度模型中实现了对数万个站点的协同预测。
Corrformer可以基于推断的天气过程传播方向生成可解释的预测，为气象科学发现洞见提供了一种完全数据驱动的人工智能范式。
Corrformer在全球、区域和城市数据集上以高置信度产生了最先进的预测结果，超越了经典统计方法和最新的深度模型，在近地面预测方面与数值方法相比也表现良好。

:triangular_flag_on_post:新闻 (2023.06) 我们的论文已在《自然机器智能》上发表，并被选为封面文章。

代码结构

|-- Corrformer
   |-- data_provider # 数据加载器
   |-- exp # 训练、验证和测试的流程
   |-- layers 
   |   |-- Embed.py # 论文中的等式(1)
   |   |-- Corrformer_EncDec.py # 论文中的等式(2)和等式(3)
   |   |-- Causal_Conv.py # 用于交叉相关的因果卷积
   |   |-- Multi_Correlation.py # 论文中的等式(5)-(10)
   |-- models
   |   |-- Corrformer.py # 整体框架
   |-- utils
   |-- scripts # 运行脚本
   |-- dataset # 存放下载的数据集
   |-- checkpoints # 存放输出或预训练模型

复现

找一台支持GPU的设备。我们的实验是在一块24GB的RTX GPU上在Linux系统中进行的。
安装Python 3.6和PyTorch 1.7.1。以下脚本可能会方便安装：

pip install -r requirements.txt # 大约需要5分钟

从[Code Ocean]下载数据集。并将它们放在./dataset文件夹下。
使用以下脚本训练和评估模型。

bash ./scripts/Global_Temp/Corrformer.sh # 大约需要18小时
bash ./scripts/Global_Wind/Corrformer.sh # 大约需要18小时

注意：由于NCEI的全球温度和全球风速原始数据已被放大10倍，这两个基准测试的实际MSE和MAE应分别除以100和10。

演示

对于简单的演示，我们建议使用预训练模型进行实验，这可以快速测试我们的代码。以下是详细说明：

按照上述说明配置环境。注意，以下实验将占用4GB GPU内存。
从[数据集]和[预训练模型]下载数据集和预训练模型。将预训练模型放在./checkpoints文件夹下。
使用以下脚本执行演示。

bash ./scripts/Demo/Global_Temp_demo.sh # 大约需要35分钟
bash ./scripts/Demo/Global_Wind_demo.sh # 大约需要35分钟

再次提醒：由于NCEI的全球温度和全球风速原始数据已被放大10倍，这两个基准测试的实际MSE和MAE应分别除以100和10。

引用

如果您觉得这个仓库有用，请引用我们的论文。

@article{wu2023corrformer,
  title={Interpretable Weather Forecasting for Worldwide Stations with a Unified Deep Model},
  author={Haixu Wu and Hang Zhou and Mingsheng Long and Jianmin Wang},
  journal={Nature Machine Intelligence},
  year={2023},
}

联系方式

如果您有任何问题或建议，欢迎联系吴海旭 (wuhaixu98@gmail.com 或 wuhx23@mails.tsinghua.edu.cn)。