作为大规模参数化偏微分方程学习替代模型的模型并行傅里叶神经算子
机器学习
2023-07-19 v3 分布式、并行与集群计算
数值分析
数值分析
摘要
傅里叶神经算子（FNOs）是近年来提出的用于学习偏微分方程（PDEs）解算子的神经网络架构，已被证明比同类深度学习方法的性能显著更优。一旦训练完成，FNOs 相比传统数值 PDE 求解器可实现多个数量级的加速。然而，由于其输入数据与网络权重的高度维数，FNOs 迄今仅应用于二维或小型三维问题。为消除这一有限问题规模的障碍，我们提出基于输入数据与网络权重域分解的模型并行版 FNO。我们证明，我们的模型并行 FNO 能够在 Perlmutter 上使用最多 512 块 A100 GPU 预测超过 26 亿变量的时变 PDE 解，并展示了在 Azure 云上训练分布式 FNO 以模拟地球 subsurface 中多相 CO 动力学的示例。
引用
@article{arxiv.2204.01205,
title = {Model-Parallel Fourier Neural Operators as Learned Surrogates for Large-Scale Parametric PDEs},
author = {Thomas J. Grady and Rishi Khan and Mathias Louboutin and Ziyi Yin and Philipp A. Witte and Ranveer Chandra and Russell J. Hewett and Felix J. Herrmann},
journal= {arXiv preprint arXiv:2204.01205},
year = {2023}
}