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.
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 oflibtorch
and rename it tolibtorch_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_cpuSimilarly, the CUDA 11.1 version of
libtorch
can be downloaded and renamed tolibtorch_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_cudaYou can also check the website, PyTorch to download the latest version of LibTorch.
Go to
ML/PYTORCH/Exec
to compile the executable. If using GPU, compile withUSE_CUDA=TRUE
. Runmake
and it should result in an executable namedmain2d.gnu.MPI.CUDA.ex
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.