The overall goal of this machine learning tutorial is to accelerate computationally expensive point-wise kernels/routines within 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.
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.
Here we use a 1-input, 2-output model to illustrate the interface between the PyTorch model and a MultiFab.
Beta Decay Reaction
In this example, the machine learning model is a regression model pre-trained to solve a two-component ODE system describing beta decay.
In the context of the pytorch model, the input is a time step
dt and output is the two-component solution of the ODE system at time
t = dt.
The TorchScript model that is included in this example is located at
If you wish to change the model, edit the
model_file parameter in
Running an AMReX application with a PyTorch model
To begin, we initialize a MultiFab full of data representing different
dt values, then copy this data into a PyTorch tensor, then call the pre-trained model to compute the outputs, and finally load the result back into a MultiFab.
The model can be evaluated on the CPU or GPU.
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 of
libtorchand rename it to
libtorch_cpu:wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.0.1%2Bcpu.zip unzip libtorch-cxx11-abi-shared-with-deps-2.0.1+cpu.zip mv libtorch libtorch_cpu
Similarly, the CUDA 11.8 version of
libtorchcan be downloaded and renamed to
libtorch_cuda:wget https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.0.1%2Bcu118.zip unzip libtorch-cxx11-abi-shared-with-deps-2.0.1+cu118.zip mv libtorch libtorch_cuda
You can also check the website, PyTorch to download the latest version of LibTorch.
ML/PYTORCH/Execto compile the executable. Run
USE_CUDA=TRUEand it should result in an executable named, e.g.,
Then you can run the example, e.g.,
mpiexec -n 4 ./main2d.gnu.MPI.ex inputs. There will be two plotfiles,
X_1at the final time).