PINNs-TF2: Fast and User-Friendly Physics-Informed Neural Networks in TensorFlow V2

TL;DR

PINNs-TF2 is a user-friendly TensorFlow V2 package that accelerates physics-informed neural networks, significantly boosting performance and simplifying implementation for solving nonlinear PDEs.

Contribution

The paper introduces PINNs-TF2, a TensorFlow V2-based package that enhances PINNs with compiler optimizations, achieving up to 119x speed-up and simplifying complex PDE modeling.

Findings

Achieved up to 119x performance boost with XLA compilers.

Demonstrated an average of 18.12x speed-up over TensorFlow V1.

Successfully handled large-scale PDE problems with many parameters.

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.