RoseNNa is a fast, portable, and minimally-intrusive library for neural network inference. It can run inference on neural networks in ONNX format, which is universal and can be used with PyTorch, TensorFlow, Keras, and more. RoseNNa's intended use case is embedding neural networks in Fortran- and C-based HPC codebases. One compiles RoseNNa and links it to an existing PDE (e.g., CFD) solver written in C or Fortran. You can then evaluate your neural network from the PDE solver at Fortran/C speeds.
RoseNNa currently supports RNNs, CNNs, and MLPs. The library is optimized Fortran and outperforms PyTorch (by a factor between 2 and 5x) for the relatively small neural networks used in physics applications, like computational fluid dynamics. RoseNNa is described in detail in A. Bati, S. H. Bryngelson (2024) Comp. Phys. Comm., 296, 109052..
program hello_roseNNa
use rosenna
implicit none
real, dimension(1,1,28,28) :: input ! model inputs
real, dimension(1,5) :: output ! model outputs
call initialize() ! reads weights
call use_model(input, output) ! run inference
end program
This example program links to the roseNNa library, parses the model inputs, and runs inference on the loaded library.
Only a few lines are required to use the library: use rosenna
, call initialize()
, and call use_model(args)
.
We have minimal dependencies. For example, on MacOS you can get away with just
brew install wget make cmake coreutils gcc
pip install torch onnx numpy fypp onnxruntime pandas
Here is a quick example of how roseNNa works. With just a few steps, you can see how to convert a basic feed-forward neural network originally built with PyTorch into usable, accurate code in Fortran.
First, cd
into the fLibrary/
directory.
Then, create PyTorch model and convert to ONNX:
python ../goldenFiles/gemm_small/gemm_small.py
Read and interpret the corresponding output files from the last step via
python modelParserONNX.py -w ../goldenFiles/gemm_small/gemm_small.onnx -f ../goldenFiles/gemm_small/gemm_small_weights.onnx
and compile the library
make library
Compile the "source files" (capiTester.f90
) and link to the library file created:
gfortran -c ../examples/capiTester.f90 -IobjFiles/
gfortran -o flibrary libcorelib.a capiTester.o
./flibrary
and finally check if the output from PyTorch model matches roseNNa's output
python ../test/testChecker.py gemm_small
-
Save the neural network model that needs to be converted
Make sure to refer to the specific library's documentation about how to save the model.
-
Convert the saved model to an ONNX format
Details on converting a saved model to ONNX format can be found on their website.
Converting an LSTM?
One important thing to note is sometimes ONNX enables optimizations that will change how the weights are stored internally (this will happen specifically for LSTMs). When converting from any library to ONNX, one should load 2 files: one with optimization and one without. This may or may not apply to all library to ONNX conversions, but here is an example using PyTorch (one with
do_constant_folding=True
and another withdo_constant_folding=False
.
#MODEL STRUCTURE FILE
torch.onnx.export(model, # model being run
(inp, hidden), # model input (or a tuple for multiple inputs)
filePath+"lstm_gemm.onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
opset_version=12, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
input_names = ['input', 'hidden_state','cell_state'], # the model's input names
output_names = ['output'], # the model's output names
)
#MODEL WEIGHTS FILE
torch.onnx.export(model, # model being run
(inp, hidden), # model input (or a tuple for multiple inputs)
filePath+"lstm_gemm_weights.onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
opset_version=12, # the ONNX version to export the model to
do_constant_folding=False, # whether to execute constant folding for optimization
input_names = ['input', 'hidden_state','cell_state'], # the model's input names
output_names = ['output'], # the model's output names
)
- Preprocess the model
fLibrary/
holds the library files that recreate and run inference on the model. Run python modelParserONNX.py -f path/to/model/structure -w path/to/weights/file
to reconstruct the model.
- Compiling the library
Then, in the same /fLibrary
directory, run make library
. This compiles the library into libcorelib.a
, which is required to link other *.o
files with the library. This library file is now ready to be integrated into any Fortran/C workflow.
One can compile a Fortran example (like the Hello RoseNNa
example above) by specifying the location of the module files and linking the library to other program files.
In practice, this looks like
gfortran -c *.f90 -Ipath/to/objFiles
gfortran -o flibrary path/to/libcorelib.a *.o
./flibrary
One can readily call roseNNa from C. Compile roseNNa, then use the following C program as an example:
void use_model(double * i0, double * o0);
void initialize(char * model_file, char * weights_file);
int main(void) {
double input[1][2] = {1,1};
double out[1][3];
initialize("onnxModel.txt","onnxWeights.txt");
use_model(input, out);
for (int i = 0; i < 3; i++) {
printf("%f ",b[0][i]);
}
}
and compile it as
gcc -c *.c
gfortran -o capi path/to/libcorelib.a *.o
./capi
Please see this document on how to extend roseNNa to new network models and this document on the details of the roseNNa pipeline.
You can cite this work as
@article{bati24,
author = {Bati, A. and Bryngelson, S. H.},
title = {{RoseNNa: A} performant, portable library for neural network inference with application to computational fluid dynamics},
journal = {Computer Physics Communications},
volume = {296},
pages = {109052},
year = {2024},
doi = {10.1016/j.cpc.2023.109052},
}