[論文レビュー] 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 は、TensorFlow 2.x と Keras を基盤に構築された完全にオープンソースで、マルチGPU対応のスケーラブルな前向きおよび逆問題向け物理情報ニューラルネットワーク（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系のトレーニング収束性を向上させること。
• トレーニング済みの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 Keras統合をシームレスに実現した完全にオープンソースでマルチGPU対応のPINNソルバーを構築できるか？
• RQ2 自己適応トレーニングは、アレン・チャン方程式のような複雑な半線形PDEの解法において、収束性と精度をどのように向上させるか？
• RQ3 スケーラビリティと複数GPUでのパフォーマンスを維持したまま、カスタムニューラルネットワークアーキテクチャをPINNソルバーに統合できる範囲はどの程度か？
• RQ4 トレーニング済みのKerasモデルをエクスポートすることで、効率的な転移学習とCPUプラットフォームでの推論のポートフォリオが可能になるか？
• RQ5 TensorDiffEqのモジュラー設計は、新しいPDEタイプ、境界条件、最適化戦略の拡張をどの程度サポートするか？

主な発見

• TensorDiffEq は、TensorFlow 2.x と Keras を用いて、PINNのスケーラブルでマルチGPU対応のトレーニングを実装し、大規模で高次元のPDEの高性能な解法を可能にした。
• フレームワークは前向き問題と逆問題の両方をサポートしており、PDEシステムにおけるパrameter推定も組み込みで対応可能で、初期条件および境界条件のサポートを備えている。
• TensorDiffEq に組み込まれた自己適応トレーニング手法は、アレン・チャン方程式のような挑戦的な半線形PDEにおいて、収束性を顕著に向上させ、標準的なコロケーション法が失敗する状況でも有効である。
• Kerasモデルのエクスポートと保存が可能であるため、転移学習やCPUプラットフォームでの推論が可能となり、デプロイの柔軟性が向上する。
• TensorDiffEq は、完全なオープンソース、マルチGPUスケーラビリティ、Keras API の拡張性、自己適応トレーニングを併せ持つ唯一の既知のPINNソルバーである。
• フレームワークは、既存のKerasツールとのシームレスな統合を示しており、研究者が複雑な物理情報付きの問題を解く際、馴染みのあるトレーニングパイプラインをそのまま利用可能にしている。