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 of
libtorchand 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
libtorchcan 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.
ML/PYTORCH/Execto compile the executable. If using GPU, compile with
makeand it should result in an executable named
Then you can run the example:
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.
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.