[논문 리뷰] PINNs-TF2: Fast and User-Friendly Physics-Informed Neural Networks in TensorFlow V2

Reza Akbarian Bafghi, Maziar Raissi|arXiv (Cornell University)|2023. 11. 07.

Model Reduction and Neural Networks인용 수 5

한 줄 요약

이 논문은 물리 기반 신경망(PINNs)의 가속을 위해 텐서플로우 2 기반으로 구축된 빠르고 사용자 친화적인 파이썬 패키지인 PINNs-TF2를 소개한다. XLA Just-In-Time(JIT) 컴파일과 정적 계산 그래프를 활용하여, 프레임워크는 텐서플로우 1.0 대비 최대 119.96배의 속도 향상을 달성했으며, 여덟 개의 PDE 예제 평균 18.12배의 향상을 보였다. 동시에 하이드라 기반의 구성 설정 워크플로우와 모듈식 조건 처리를 통해 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의 구현을 단순화하기 위해.
• 하이드라 프레임워크를 활용한 구성 가능한 워크플로우를 통해 복잡한 PDE 설정을 추상화하여 사용자 경험을 향상시키기 위해.
• 컴파일러, 배치 크기, 학습 가능한 파라미터의 영향을 PINN 훈련의 효율성과 정확도에 대해 평가하기 위해.
• 성능과 사용성 향상을 위해 공개된 오픈소스 도구를 제공하여 PINN 연구에 기여하기 위해.

제안 방법

• 패키지는 tf.function을 사용하여 텐서플로우 2의 정적 계산 그래프를 활용하고, jit_compile=True를 설정하여 최적화된 기계어 생성을 위한 XLA JIT 컴파일을 활성화한다.
• 하이드라 프레임워크를 통합하여 구성 파일을 관리하고, 공간/시간 도메인, 샘플링 전략, 경계 조건, 신경망 아키텍처를 정의한다.
• 프레임워크는 연속적 및 이산적 시간 모델을 모두 지원하며, 콜로케이션 포인트와 해의 데이터를 사용하여 PDE 및 경계 조건의 손실 함수를 정의한다.
• 데이터, 콜로케이션, 경계 조건에 대해 SSE 및 MSE와 같은 다중 손실 함수를 하나의 신경망에 공유하여 조합한다.
• 훈련 및 평가 단계는 XLA를 통해 컴iles되며, 이는 연산을 융합하고 대상 하드웨어에서 메모리와 계산을 최적화한다.
• 프레임워크는 정방향 및 역방향 문제 모두를 지원하며, 버거스 및 라우-스톡스 방정식과 같은 매개변수형 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 XLA JIT 컴파일은 표준 텐서플로우 1.0 및 컴파일 없이 사용된 TF2 대비 PINN 훈련 속도에 어떤 영향을 미치는가?
• RQ2 하이드라 기반의 모듈식 구성 기반 프레임워크인 PINNs-TF2는 다양한 PDE 문제의 PINN 구현을 얼마나 단순화할 수 있는가?
• RQ3 배치 크기와 학습 가능한 파라미터 수는 복잡한 PDE에서 PINN의 성능과 수렴에 어떤 영향을 미치는가?
• RQ4 실제 사례 연구를 통해 보여지듯이, 많은 파라미터와 높은 배치 크기를 가진 대규모 문제를 효율적으로 처리할 수 있는가?
• RQ5 JIT 컴파일을 사용할 경우 속도 향상과 수치적 안정성 사이의 상충 관계는 어떠한가?

주요 결과

• PINNs-TF2는 원래의 텐서플로우 1.0 대비 최대 119.96배의 속도 향상을 달성했으며, 주로 XLA JIT 컴파일 덕분이었다.
• 여덟 가지 다양한 PDE 예제 평균으로 18.12배의 속도 향상을 보였다. 이는 버거스 및 라우-스톡스 방정식 포함.
• 이산 정방향 버거스 방정식의 경우, JIT 컴파일만으로도 원래 TF1 대비 11.39배의 속도 향상을 달성했다.
• 이산 역방향 버거스 방정식의 경우, JIT 컴파일로 15.77배의 속도 향상을 얻었으며, 해의 정확도에 대한 영향은 최소한이었다 (평균 상대 오차 0.142).
• 하이드라 프레임워크 통합으로 PDE 문제의 효율적이고 모듈식인 구성이 가능해져 재현성과 사용자 접근성 향상이 이루어졌다.
• 3D 라우-스톡스 방정식에 대해 10,000의 배치 크기와 10×250층의 네트워크를 사용한 대규모 문제에서도 스케일링 가능성을 입증하여 실용적 타당성을 확인했다.

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.