RoseNNa: A performant, portable library for neural network inference with application to computational fluid dynamics

Abstract

The rise of neural network-based machine learning ushered in high-level libraries, including TensorFlow and PyTorch, to support their functionality. Computational fluid dynamics (CFD) researchers have benefited from this trend and produced powerful neural networks that promise shorter simulation times. For example, multilayer perceptrons (MLPs) and Long Short Term Memory (LSTM) recurrent-based (RNN) architectures can represent sub-grid physical effects, like turbulence. Implementing neural networks in CFD solvers is challenging because the programming languages used for machine learning and CFD are mostly non-overlapping, We present the roseNNa library, which bridges the gap between neural network inference and CFD. RoseNNa is a non-invasive, lightweight (1000 lines), and performant tool for neural network inference, with focus on the smaller networks used to augment PDE solvers, like those of CFD, which are typically written in C/C++ or Fortran. RoseNNa accomplishes this by automatically converting trained models from typical neural network training packages into a high-performance Fortran library with C and Fortran APIs. This reduces the effort needed to access trained neural networks and maintains performance in the PDE solvers that CFD researchers build and rely upon. Results show that RoseNNa reliably outperforms PyTorch (Python) and libtorch (C++) on MLPs and LSTM RNNs with less than 100 hidden layers and 100 neurons per layer, even after removing the overhead cost of API calls. Speedups range from a factor of about 10 and 2 faster than these established libraries for the smaller and larger ends of the neural network size ranges tested. Program summaryProgram Title: RoseNNaCPC Library link to program files:https://doi.org/10.17632/srbrfx8k74.1Developer's repository link:https://github.com/comp-physics/roseNNaLicensing provisions: MITProgramming language: Fortran90 and PythonNature of problem: Neural network applications in computational fluid dynamics exhibit high promise. However, they are often not deployed on HPC systems due to limited deep-learning support in Fortran or C/C++. Previous attempts to solve this issue rely on significant manual intervention. This leaves researchers with no practical solution for conveniently porting their trained models for use in performant HPC codebases.Solution method: RoseNNa converts trained neural networks from popular machine learning frameworks to Fortran/C code for inference, requiring minimal manual intervention. Metaprogramming enables automatic code creation, and ONNX ensures flexible machine learning library support.