ML/PYTORCH

The overall goal of machine learning models in this context is to accelerate computationally expensive kernels/routines as part of an AMReX simulation. This tutorial demonstrates how to interface a pre-trained PyTorch machine learning model to an AMReX simulation by querying inputs from and supplying outputs to an AMReX MultiFab. Here we use a 1-input, 2-output model to illustrate the interface between the PyTorch model and a MultiFab.

PyTorch is a commonly used machine learning package with a C++ API library called LibTorch. Located in the directory amrex-tutorials/ExampleCodes/ML/PYTORCH, this example uses a machine learning model to solve a radioactive beta decay problem. To begin, we initialize data on a MultiFab, then copy the data into a PyTorch tensor, then we call the pre-trained model to compute the outputs, and finally we load the result back into a MultiFab. The program runs on either only the CPU or both the CPU and GPU.

Running an AMReX application with a PyTorch model

Below is a step-by-step guide to successfully run an AMReX program that uses a PyTorch model. It will require the model to have been saved as a TorchScript. In this example the TorchScript file is model.pt. For more information on TorchScript, please see their intro tutorial.

1. Before compiling, either a CPU or CUDA version of LibTorch (PyTorch C++ library) must be downloaded into ML/PYTORCH/. To download the CPU-only version of libtorch and rename it to libtorch_cpu:

   wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-1.9.0%2Bcpu.zip
   unzip libtorch-cxx11-abi-shared-with-deps-1.9.0+cpu.zip
   mv libtorch libtorch_cpu
   

   Similarly, the CUDA 11.1 version of libtorch can be downloaded and renamed to libtorch_cuda:

   wget https://download.pytorch.org/libtorch/cu111/libtorch-cxx11-abi-shared-with-deps-1.9.0%2Bcu111.zip
   unzip libtorch-cxx11-abi-shared-with-deps-1.9.0+cu111.zip
   mv libtorch libtorch_cuda
   

   You can also check the website, PyTorch to download the latest version of LibTorch.

2. Go to ML/PYTORCH/Exec to compile the executable. If using GPU, compile with USE_CUDA=TRUE. Run make and it should result in an executable named main2d.gnu.MPI.CUDA.ex

3. Then you can run the example: ./main2d.gnu.MPI.CUDA.ex inputs.

Beta Decay

In this example, the machine learning model is a regression model pre-trained to solve a two-component ODE system describing beta decay. The input is a time step dt and output is the two-component solution of the ODE system at time t = dt.

Pre-trained Model

The TorchScript model that is included in this example is located at ML/PYTORCH/Exec/model.pt. If you wish to change the model, edit the model_file parameter in ML/PYTORCH/Exec/inputs to your desired PyTorch model file location.