[论文解读] PINNs-TF2: Fast and User-Friendly Physics-Informed Neural Networks in TensorFlow V2

Reza Akbarian Bafghi, Maziar Raissi|arXiv (Cornell University)|Nov 7, 2023

Model Reduction and Neural Networks被引用 5

一句话总结

本论文介绍了 PINNs-TF2，这是一个基于 TensorFlow 2 的快速且用户友好的 Python 包，旨在加速物理信息神经网络（PINNs）的训练。通过利用 XLA 即时（JIT）编译和静态计算图，该框架在八个偏微分方程（PDE）示例中实现了相较于 TensorFlow 1.0 的最高 119.96 倍加速，平均提升达 18.12 倍，同时通过 Hydra 配置工作流和模块化条件处理，显著简化了 PINN 的实现。

ABSTRACT

Physics-informed neural networks (PINNs) have gained prominence for their capability to tackle supervised learning tasks that conform to physical laws, notably nonlinear partial differential equations (PDEs). This paper presents "PINNs-TF2", a Python package built on the TensorFlow V2 framework. It not only accelerates PINNs implementation but also simplifies user interactions by abstracting complex PDE challenges. We underscore the pivotal role of compilers in PINNs, highlighting their ability to boost performance by up to 119x. Across eight diverse examples, our package, integrated with XLA compilers, demonstrated its flexibility and achieved an average speed-up of 18.12 times over TensorFlow V1. Moreover, a real-world case study is implemented to underscore the compilers' potential to handle many trainable parameters and large batch sizes. For community engagement and future enhancements, our package's source code is openly available at: https://github.com/rezaakb/pinns-tf2.

研究动机与目标

• 解决由于动态计算图中重复梯度计算导致的 PINN 训练性能瓶颈。
• 简化涉及非线性 PDE 的正向与反向问题中 PINN 的实现。
• 通过使用 Hydra 框架的可配置工作流，抽象复杂 PDE 设置，提升用户体验。
• 评估编译器、批量大小和可训练参数数量对 PINN 训练效率和准确率的影响。
• 提供一个公开可用的开源工具，以提升 PINN 研究的性能与可用性。

提出的方法

• 该包通过 tf.function 并设置 jit_compile=True，利用 TensorFlow 2 的静态计算图实现 XLA JIT 编译，以生成优化的机器代码。
• 集成 Hydra 框架以管理配置文件，定义空间/时间域、采样策略、边界条件和神经网络架构。
• 该框架支持连续时间和离散时间模型，通过配点和解数据定义 PDE 及边界条件的损失函数。
• 多个损失函数（如 SSE 和 MSE）被组合使用，用于数据、配点和边界条件，并在单个神经网络中共享。
• 训练和评估步骤通过 XLA 编译，实现操作融合，并在目标硬件上优化内存与计算。
• 该方法支持正向与反向问题，包括 Burgers’ 方程和 Navier-Stokes 方程等参数化 PDE，支持可配置的超参数与批量大小。

Figure 1: Overview of the simplified PINNs-TF2 framework: (a) Users provide a config file and define reading data and PDE functions. The package then processes them to initialize a neural network, formulate a mesh for the data, and establish conditions. (b) Elucidates the components of the condition

实验结果

研究问题

• RQ1 与标准的 TensorFlow 1.0 及未启用编译的 TF2 相比，XLA JIT 编译对 PINN 训练速度有何影响？
• RQ2 像 PINNs-TF2 这类模块化、配置驱动的框架，在多大程度上能简化 PINN 中多样化 PDE 问题的实现？
• RQ3 批量大小和可训练参数数量如何影响复杂 PDE 中 PINN 的性能与收敛性？
• RQ4 该框架能否高效处理大规模问题（如含大量参数和高批量大小的问题），如实际案例研究所示？
• RQ5 在 PINN 训练中使用 JIT 编译时，加速与数值稳定性之间的权衡如何？

主要发现

• PINNs-TF2 在原始 TensorFlow 1.0 实现上实现了最高 119.96 倍的加速，主要得益于 XLA JIT 编译。
• 在八个不同的 PDE 示例中，该框架平均实现了 18.12 倍的加速，包括 Burgers’ 方程和 Navier-Stokes 方程。
• 对于离散正向 Burgers’ 方程，仅使用 JIT 编译即实现了相较于原始 TF1 实现 11.39 倍的加速。
• 在离散反向 Burgers’ 方程中，JIT 编译带来了 15.77 倍的加速，且解的精度损失极小（平均相对误差为 0.142）。
• Hydra 框架的集成实现了 PDE 问题的高效、模块化配置，提升了可复现性与用户可访问性。
• 该包在大规模问题上表现出良好的可扩展性，例如在 10,000 批量大小和 10×250 层网络的 3D Navier-Stokes 方程中，证实了其实际可行性。

Figure 2: Each subplot denotes a unique problem, with its specific iteration count indicated at the top. The logarithmic x-axis shows speed-up relative to TF1, and the y-axis the mean error, highlighting JIT compiler boosts speed without added error.