/pinn-heat-equation

A standalone project to test libtorch C++ APIs on solving the 2D heat equation with PINN.

Primary LanguageC++

Solving 2D Heat Equation with PINN

This project is to solve a 2D heat equation with PINN. The math description and Python implementation is given by the Jupyter script. It is one of the JLab EPSCI PHASM examples.

This stand-alone repo is created to test the libtorch C++ APIs without considering the compatibility to the other PHASM codebase. It can also be used as the backup of the main repo.

Environment

The repo is tested on the JLab ifarm GPU node with below configurations.

glibc(ldd): 2.17
cmake: 3.21.1
gcc: 10.2.0
CUDA: 11.4
GPU: one Tesla T4 or one TitanRTX
cuDNN: 8.4.1
libtorch: 1.12.1 + Linux + LibTorch + C++/Java + cu11.3 + pre-cxx11abi
OS: CentOS Linux 7 (Core)
Kernel: Linux 3.10.0-1160.71.1.el7.x86_64

As I have no root control of the system, the library version problem is not well handled. cxx11 ABI libtorch should be used but failed with a low glibc (description here) problem. Then I use the pre-cxx11 ABI libtorch for now.

Build the project

# add torch lib to path
export LD_LIBRARY_PATH=/path/to/libtorch/lib:$LD_LIBRARY_PATH

# build project
mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH="/path/to/libtorch;/path/to/cuDNN" ..
make

libtorch and cuDNN installation

cuDNN and libtorch can be accessed at the official sites. The packages I am using are given as below.

cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive.tar.xz
libtorch-cxx11-abi-shared-with-deps-1.12.1+cu113.zip  # cxx11 ABI, not used
libtorch-shared-with-deps-1.12.0+cu113.zip  # Pre-cxx11 ABI, used in this repo

Unzip these packages, and they will make up of your /path/to/libtorch and /path/to/cuDNN.

NCU profiling

A ncu profiling is conducted on the training process to study the behaviours of the kernels. The Roofline model is utilized to identify the bottleneck of the application.

TODOs

  • Add the original Gauss-Seidel iteration and timing functions for the whole process.
  • Reshape the NN structure to achieve higher throuput.

References