Native Fortran Implementation of TensorFlow-Trained Deep and Bayesian Neural Networks
Aidan Furlong, Xingang Zhao, Bob Salko, Xu Wu
TL;DR
This work tackles the challenge of running TensorFlow-trained ML models directly in Fortran-based nuclear engineering codes by introducing a lightweight, native framework that reads HDF5-exported DNN and BNN configurations. The approach centralizes model loading, data standardization, and prediction in Fortran modules, enabling high-throughput inference without external runtimes or ONNX conversion. Verification across a noisy sinusoid and a nuclear-relevant CHF surrogate demonstrated strong agreement with TensorFlow while delivering substantial speedups (up to ~19.6x for DNN and ~8x for BNN) and practical UQ capabilities. The framework supports iterative solvers and ensembles for UQ, with successful integration considerations for the CTF code, highlighting its potential to accelerate ML-assisted nuclear simulations in production environments.
Abstract
Over the past decade, the investigation of machine learning (ML) within the field of nuclear engineering has grown significantly. With many approaches reaching maturity, the next phase of investigation will determine the feasibility and usefulness of ML model implementation in a production setting. Several of the codes used for reactor design and assessment are primarily written in the Fortran language, which is not immediately compatible with TensorFlow-trained ML models. This study presents a framework for implementing deep neural networks (DNNs) and Bayesian neural networks (BNNs) in Fortran, allowing for native execution without TensorFlow's C API, Python runtime, or ONNX conversion. Designed for ease of use and computational efficiency, the framework can be implemented in any Fortran code, supporting iterative solvers and UQ via ensembles or BNNs. Verification was performed using a two-input, one-output test case composed of a noisy sinusoid to compare Fortran-based predictions to those from TensorFlow. The DNN predictions showed negligible differences and achieved a 19.6x speedup, whereas the BNN predictions exhibited minor disagreement, plausibly due to differences in random number generation. An 8.0x speedup was noted for BNN inference. The approach was then further verified on a nuclear-relevant problem predicting critical heat flux (CHF), which demonstrated similar behavior along with significant computational gains. Discussion regarding the framework's successful integration into the CTF thermal-hydraulics code is also included, outlining its practical usefulness. Overall, this framework was shown to be effective at implementing both DNN and BNN model inference within Fortran, allowing for the continued study of ML-based methods in real-world nuclear applications.