[論文レビュー] 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

ひとこと要約

この論文では、TensorFlow 2に基づいて開発された、物理に基づくニューラルネットワーク（PINNs）の高速化を図る、高速で使いやすいPythonパッケージであるPINNs-TF2を紹介する。XLAのジャストインタイム（JIT）コンパイルと静的計算グラフを活用することで、TensorFlow 1.0と比較して最大119.96倍の高速化を達成し、8つのPDE例において平均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を用いてtf.functionにjit_compile=Trueを指定することで、TensorFlow 2の静的計算グラフを活用し、最適化されたマシンコード生成のためのXLA JITコンパイルを有効化する。
• Hydraフレームワークを統合し、空間的・時間的領域、サンプリング戦略、境界条件、ニューラルネットワークアーキテクチャを定義する構成ファイルを管理する。
• 本フレームワークは連続時刻および離散時刻モデルをサポートし、コロケーションポイントと解のデータを用いて、PDEおよび境界条件の損失関数を定義する。
• データ、コロケーション、境界条件のための複数の損失関数（例：SSEとMSE）を、単一のニューラルネットワークに共有して組み合わせる。
• 学習および評価ステップはXLAでコンパイルされ、演算の統合と、ターゲットハードウェア上のメモリおよび計算最適化が行われる。
• 本手法は前向きおよび逆問題をサポートし、Burgers’方程式やナビエ-ストークス方程式などのパラメトリックPDEに対しても、構成可能なハイパーパrameterとバッチサイズを備える。

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 XLA JITコンパイルは、標準のTensorFlow 1.0およびコンパイルなしのTF2と比較して、PINNの学習速度にどのように影響するか？
• RQ2 PINNs-TF2のようなモジュラーで構成駆動のフレームワークは、PINNにおける多様なPDE問題の実装をどの程度簡素化できるか？
• RQ3 バッチサイズとトレーニング可能なパラメータの数は、複雑なPDEにおけるPINNのパフォーマンスと収束にどのように影響するか？
• RQ4 本フレームワークは、多くのパラメータと高いバッチサイズを有する大規模問題を、実世界の事例で示されるように効率的に処理できるか？
• RQ5 JITコンパイルを用いる際の、高速化と数値的安定性のトレードオフは何か？

主な発見

• PINNs-TF2は、元のTensorFlow 1.0実装と比較して最大119.96倍の高速化を達成した。主な要因はXLA JITコンパイルであった。
• 8つの多様なPDE例（Burgers’方程式やナビエ-ストークス方程式を含む）において、平均して18.12倍の高速化を達成した。
• 離散的前向きBurgers’方程式では、JITコンパイルのみで元のTF1実装と比較して11.39倍の高速化が達成された。
• 離散的逆問題Burgers’方程式では、JITコンパイルにより15.77倍の高速化が得られ、解の精度にほとんど劣化がなかった（平均相対誤差0.142）。
• Hydraフレームワークの統合により、PDE問題の効率的でモジュラーな構成が可能になり、再現性とユーザーのアクセス性が向上した。
• 本パッケージは、10,000のバッチサイズと10×250層ネットワークを有する3次元ナビエ-ストークス方程式のような大規模問題においてもスケーラビリティを示し、実用的妥当性を確認した。

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.