[论文解读] TensorDiffEq: Scalable Multi-GPU Forward and Inverse Solvers for Physics Informed Neural Networks

Levi D. McClenny, Mulugeta Haile+1 more|arXiv (Cornell University)|Mar 30, 2021

Model Reduction and Neural Networks参考文献 22被引用 11

一句话总结

TensorDiffEq 是一个完全开源、支持多 GPU 的框架，基于 TensorFlow 2.x 和 Keras 构建，用于可扩展的前向和反向物理信息神经网络（PINN）求解器。它通过类似 Keras 的模型定义方式支持高维偏微分方程（PDE）求解，支持自定义网络架构，并引入自适应训练以提升收敛性，为科学研究中的机器学习提供一个生产就绪、模块化的替代方案，相较于单 GPU 或闭源的 PINN 工具更具优势。

ABSTRACT

Physics-Informed Neural Networks promise to revolutionize science and engineering practice, by introducing domain-aware deep machine learning models into scientific computation. Several software suites have emerged to make the implementation and usage of these architectures available to the research and industry communities. Here we introduce\linebreak TensorDiffEq, built on Tensorflow 2.x, which presents an intuitive Keras-like interface for problem domain definition, model definition, and solution of forward and inverse problems using physics-aware deep learning methods. TensorDiffEq takes full advantage of Tensorflow 2.x infrastructure for deployment on multiple GPUs, allowing the implementation of large high-dimensional and complex models. Simultaneously, TensorDiffEq supports the Keras API for custom neural network architecture definitions. In the case of smaller or simpler models, the package allows for rapid deployment on smaller-scale CPU platforms with negligible changes to the implementation scripts. We demonstrate the basic usage and capabilities of TensorDiffEq in solving forward, inverse, and data assimilation problems of varying sizes and levels of complexity. The source code is available at https://github.com/tensordiffeq.

研究动机与目标

• 解决科学机器学习中缺乏完全开源、可扩展的多 GPU PINN 求解器的问题。
• 提供一个模块化、可扩展的框架，支持使用物理信息深度学习求解前向和反向问题。
• 使研究人员能够通过 Keras API 定义自定义神经网络架构，同时在多 GPU 上保持高性能训练。
• 通过自适应训练方法提升复杂 PDE 系统的训练收敛性，特别是对像 Allen-Cahn 这类半线性方程的求解。
• 通过支持导出和重用训练后的 Keras 模型，提升部署的可移植性，实现在 CPU 和 GPU 平台间的跨平台使用。

提出的方法

• 该框架使用类似 Keras 的接口定义神经网络架构，实现与 Keras 生态系统的完整集成，用于模型构建和优化。
• 利用 TensorFlow 2.x 的自动微分功能，计算神经网络输出在空间和时间维度上的导数，用于 PDE 残差计算。
• 残差损失通过物理信息约束计算：r(x,t;w) = N_x,t[u(x,t;w)]，其中 u 为网络输出，N 为微分算子。
• 损失函数结合了数据损失、残差损失、边界条件损失和初始条件损失，支持前向和反向问题的求解。
• 通过在优化过程中动态调整损失分量的权重，实现自适应训练，以提升收敛稳定性。
• 通过 TensorFlow 的分布式策略支持多 GPU 训练，实现对大规模、高维 PDE 问题的可扩展训练。

实验结果

研究问题

• RQ1 能否构建一个完全开源、支持多 GPU 的 PINN 求解器，实现与 Keras 的无缝集成，以支持前向和反向问题？
• RQ2 自适应训练在求解复杂半线性 PDE（如 Allen-Cahn 方程）时，如何提升收敛性和精度？
• RQ3 在保持多 GPU 上的可扩展性和性能的前提下，自定义神经网络架构在 PINN 求解器中的集成程度如何？
• RQ4 该框架是否能通过导出训练后的 Keras 模型，实现高效的迁移学习和推理可移植性？
• RQ5 TensorDiffEq 的模块化设计在支持新类型 PDE、边界条件和优化策略方面，其可扩展性如何体现？

主要发现

• TensorDiffEq 通过 TensorFlow 2.x 和 Keras 成功实现了 PINN 的可扩展多 GPU 训练，能够高效求解大规模、高维的 PDE 问题。
• 该框架支持前向和反向问题，包括 PDE 系统中的参数估计，并内置了对初始条件和边界条件的支持。
• TensorDiffEq 中的自适应训练方法显著提升了对复杂半线性 PDE（如 Allen-Cahn 方程）的收敛性，而标准配点法在此类问题上会失效。
• 支持导出和保存 Keras 模型，使得在 CPU 平台上进行迁移学习和推理成为可能，显著增强了部署灵活性。
• TensorDiffEq 是目前已知唯一同时具备完全开源、多 GPU 可扩展性、Keras API 可扩展性以及自适应训练能力的 PINN 求解器。
• 该框架与现有 Keras 工具实现了无缝集成，使研究人员能够在求解复杂物理信息问题时，沿用熟悉的训练流程。