Reducing the Overhead of Coupled Machine Learning Models between Python and Fortran

RMetS Early Career Conference, Reading

Jack Atkinson

ICCS/Cambridge

Simon Clifford

ICCS/Cambridge

Athena Elafrou

NVIDIA

Minah Yang

NYU

Dave Connelly

NYU

+ Others

2024-06-09

Machine Learning in geophysical sciences

Machine learning is everywhere

Some applications in geosciences: ¹

• Emulation of existing parameterisations
  (Espinosa et al. 2022)
  
  
• Data-driven paramterisations
  (Yuval and O’Gorman 2020; Giglio, Lyubchich, and Mazloff 2018)
  
  
• Downscaling/Upsampling
  (Harris et al. 2022)
  

• Time series forecasting
  (Shao et al. 2021)
  
  
• Equation discovery
  (Zanna and Bolton 2020; Ma et al. 2021)
  
  
  
• Complete forecasting
  (Rasp et al. 2020; Pathak et al. 2022; Bi et al. 2022)
  
  

Speaker notes go here.

1. See review of Kashinath et al. (2021)

Challenges

Replacing physics-based components of larger models (emulation or data-driven).

• Language interoperation
• Physical compatibility

Mathematical Bridge by cmglee used under CC BY-SA 3.0

Emulation and data-driven is a focus of several of out projects.

Models typically in Fortran ML typically in Python Challenge of interoperation

Solutions

Possible solutions

Ideally need to:

• Not generate excess additional work for user
  • Not require excess knowledge of computing
  • Minimal learning curve
• Not add excess dependencies
• Be easy to maintain
• Maximise performance

1. No re-writing net after you have already done and trained
2. For scientists doing science, not compsci
3. Run on HPC environments so minimal additional
4. Easily keep up to date
5. Simple to learn and deploy
6. Needs to be as efficient as possible

Possible solutions

• Implement a NN in Fortran
• Forpy/CFFI
• SmartSim/Pipes
• Fortran-Keras Bridge

Possible solutions

• Implement a NN in Fortran
  • e.g. inference-engine, neural-fortran, own solution etc.
  • how do you ensure you port it correctly?
  • ML libraries are highly optimised, probably more so than your code.
• Forpy/CFFI
• SmartSim/Pipes
• Fortran-Keras Bridge

Fortran: - Existing: - more dependencies - new skillset required - longevity/featureset? - Self written: - more to maintain - needs checks and tests - need to be careful you exactly match the implementation

Possible solutions

• Implement a NN in Fortran
• Forpy/CFFI
  • easy to add forpy.mod file and compile
  • need to manage and link python environment
  • increases dependencies
• SmartSim/Pipes
• Fortran-Keras Bridge

Possible solutions

• Implement a NN in Fortran
• Forpy/CFFI
• SmartSim/Pipes
  • pass data between workers through a network glue layer
  • steep learning curve
  • involves data copying
• Fortran-Keras Bridge

Possible solutions

• Implement a NN in Fortran
• Forpy/CFFI
• SmartSim/Pipes
• Fortran-Keras Bridge
  • pure Fortran
  • inactive and incomplete

Our Solution

Python
env

Python
runtime

xkcd #1987 by Randall Munroe, used under CC BY-NC 2.5

Our Solution

Ftorch and TF-lib

• Use Fortran’s intrinsic C-bindings to access the C/C++ APIs provided by ML frameworks
• Performance
• Ease of use
• Use frameworks’ implementations directly

MOTIVATION Driven by disadvantages of the above. Make it quick and easy for researchers to correctly deploy their models.

• intrinsic => readily available
• highly functional

Our Solution

Ftorch and TF-lib

• Use Fortran’s intrinsic C-bindings to access the C/C++ APIs provided by ML frameworks
• Performance
  • avoids python runtime
  • no-copy transfer (shared memory)
• Ease of use
• Use frameworks’ implementations directly

Copying is one of the most expensive operations you can do. Avoid this!

Our Solution

Ftorch and TF-lib

• Use Fortran’s intrinsic C-bindings to access the C/C++ APIs provided by ML frameworks
• Performance
• Ease of use
  • pleasant API (see next slides)
  • utilities for generating TorchScript/TF module provided
  • examples provided
• Use frameworks’ implementations directly

Our Solution

Ftorch and TF-lib

• Use Fortran’s intrinsic C-bindings to access the C/C++ APIs provided by ML frameworks
• Performance
• Ease of use
• Use frameworks’ implementations directly
  • feature support
  • future support
  • direct translation of python models ¹

• NB Currently inference only - build and train in python

1. Caveats exist if using exotic non-Torch features.

Code Example - PyTorch

• Take model file
• Save as torchscript

import torch
import my_ml_model

trained_model = my_ml_model.initialize()
scripted_model = torch.jit.script(model)
scripted_model.save("my_torchscript_model.pt")

Code Example - PyTorch

Neccessary imports:

use, intrinsic :: iso_c_binding, only: c_int64_t, c_float, c_char, &
                                       c_null_char, c_ptr, c_loc
use ftorch

Loading a pytorch model:

model = torch_module_load('/path/to/saved/model.pt'//c_null_char)

Code Example - PyTorch

Tensor creation from Fortran arrays:

! Fortran variables
real, dimension(:,:), target  :: SST, model_output
! C/Torch variables
integer(c_int), parameter :: dims_T = 2
integer(c_int64_t) :: shape_T(dims_T)
integer(c_int), parameter :: n_inputs = 1
type(torch_tensor), dimension(n_inputs), target :: model_inputs
type(torch_tensor) :: model_output_T

shape_T = shape(SST)

model_inputs(1) = torch_tensor_from_blob(c_loc(SST), dims_T, shape_T &
                                         torch_kFloat64, torch_kCPU)

model_output = torch_tensor_from_blob(c_loc(output), dims_T, shape_T, &
                                      torch_kFloat64, torch_kCPU)

Code Example - PyTorch

Running the model

call torch_module_forward(model, model_inputs, n_inputs, model_output_T)

Cleaning up:

call torch_tensor_delete(model_inputs(1))
call torch_module_delete(model)

We recommend PT TF has a weakness that prior knowledge of model structure is required.

Application

Gravity Wave parameterisation in MiMA

• Model of an idealised Moist Atmosphere (Jucker and Gerber 2017)
  • Aquaplanet model with moisture and mixed-layer ocean
  • Successful implementation of gravity waves produces QBO ¹

• Neural Net
  • Emulate Alexander and Dunkerton (1999) gravity wave parameterisation.
  • Fully-connected multi-layer net with identical Pytorch and TensorFlow versions
  • Initially interfaced (slowly) using Forpy library in Espinosa et al. (2022)

Too slow by a factor of 10+ - led to question of why are we doing this!

1. Quasi-biennial oscillation

Tests

Investigate the coupling between Fortran, and the PyTorch and Tensorflow models.

• Replace AD99 GW parameterisation in MiMA
• Run MiMA for 3 years¹
• Time just the GW parameterisation
• Compare online simulation results to an AD99 run
  • Tensorflow and PyTorch
  • Old forpy and our new couplings

1. to capture one QBO

Results

Time-height sections of zonal-mean zonal wind (m/s) averaged between ±5^o longitude

Timings (real seconds) for computing gravity wave drag.¹

	Forpy	Direct	% Direct/Forpy
PyTorch	94.43 s	134.81 s	142.8 %
TensorFlow	667.16 s	170.31 s	25.5 %

1. For reference AD99 takes 69.25 s.

Conclusions

Take away messages

• Machine learning has many potential applications to enhance NWP and ESM
• Leveraging it effectively requires care
• We have developed libraries to allow scientists to easily and efficiently deploy ML within Fortran models
• For new projects we advise using PyTorch

Software

• The libraries presented can be found at:
  • Torch: https://github.com/Cambridge-ICCS/fortran-pytorch-lib
  • TensorFlow: https://github.com/Cambridge-ICCS/fortran-tf-lib
  • MIT licensed
  • Documentation provided
• Their implementation in the MiMA model can be found at
  https://github.com/DataWaveProject/MiMA-machine-learning

Future work

• Provide a tagged first release on github
  • Publication through JOSS and zenodo
• Further test GPU functionalities
• Implement functionalities beyond inference?
  • Online training likely to become important
• Further testing by external users
  • Help and feedback wanted!

References

Alexander, MJ, and TJ Dunkerton. 1999. “A Spectral Parameterization of Mean-Flow Forcing Due to Breaking Gravity Waves.” Journal of the Atmospheric Sciences 56 (24): 4167–82.

Bi, Kaifeng, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian. 2022. “Pangu-Weather: A 3d High-Resolution Model for Fast and Accurate Global Weather Forecast.” arXiv Preprint arXiv:2211.02556.

Espinosa, Zachary I, Aditi Sheshadri, Gerald R Cain, Edwin P Gerber, and Kevin J DallaSanta. 2022. “Machine Learning Gravity Wave Parameterization Generalizes to Capture the QBO and Response to Increased CO2.” Geophysical Research Letters 49 (8): e2022GL098174.

Giglio, Donata, Vyacheslav Lyubchich, and Matthew R Mazloff. 2018. “Estimating Oxygen in the Southern Ocean Using Argo Temperature and Salinity.” Journal of Geophysical Research: Oceans 123 (6): 4280–97.

Harris, Lucy, Andrew TT McRae, Matthew Chantry, Peter D Dueben, and Tim N Palmer. 2022. “A Generative Deep Learning Approach to Stochastic Downscaling of Precipitation Forecasts.” Journal of Advances in Modeling Earth Systems 14 (10): e2022MS003120.

Jucker, Martin, and EP Gerber. 2017. “Untangling the Annual Cycle of the Tropical Tropopause Layer with an Idealized Moist Model.” Journal of Climate 30 (18): 7339–58.

Kashinath, Karthik, M Mustafa, Adrian Albert, JL Wu, C Jiang, Soheil Esmaeilzadeh, Kamyar Azizzadenesheli, et al. 2021. “Physics-Informed Machine Learning: Case Studies for Weather and Climate Modelling.” Philosophical Transactions of the Royal Society A 379 (2194): 20200093.

Ma, Donglai, Jacob Bortnik, Edurado Alves, Enrico Camporeale, Xiangning Chu, and Adam Kellerman. 2021. “Data-Driven Discovery of the Governing Equations Describing Radiation Belt Dynamics.” In AGU Fall Meeting Abstracts, 2021:SA15B–1928.

Pathak, Jaideep, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, et al. 2022. “Fourcastnet: A Global Data-Driven High-Resolution Weather Model Using Adaptive Fourier Neural Operators.” arXiv Preprint arXiv:2202.11214.

Rasp, Stephan, Peter D Dueben, Sebastian Scher, Jonathan A Weyn, Soukayna Mouatadid, and Nils Thuerey. 2020. “WeatherBench: A Benchmark Data Set for Data-Driven Weather Forecasting.” Journal of Advances in Modeling Earth Systems 12 (11): e2020MS002203.

Shao, Qi, Wei Li, Guijun Han, Guangchao Hou, Siyuan Liu, Yantian Gong, and Ping Qu. 2021. “A Deep Learning Model for Forecasting Sea Surface Height Anomalies and Temperatures in the South China Sea.” Journal of Geophysical Research: Oceans 126 (7): e2021JC017515.

Yuval, Janni, and Paul A O’Gorman. 2020. “Stable Machine-Learning Parameterization of Subgrid Processes for Climate Modeling at a Range of Resolutions.” Nature Communications 11 (1): 3295.

Zanna, Laure, and Thomas Bolton. 2020. “Data-Driven Equation Discovery of Ocean Mesoscale Closures.” Geophysical Research Letters 47 (17): e2020GL088376.

Closing slide, thanks, and questions

Libraries can be found at:

• Torch: https://github.com/Cambridge-ICCS/fortran-pytorch-lib
• TensorFlow: https://github.com/Cambridge-ICCS/fortran-tf-lib

Thanks to:

Get in touch:

 Jack Atkinson

 jackatkinson.net

 jwa34[AT]cam.ac.uk

 jatkinson1000

 @jatkinson1000@fosstodon.org

Slides available at: https://jackatkinson.net/slides/RMetS/RMetS.html