ML

This topic shows how to evaluate machine-learning models on AMReX data. The tutorial needs LibTorch, the C++ API of PyTorch, and is built only with -DAMReXTutorials_TORCH=ON, with CMAKE_PREFIX_PATH or Torch_DIR pointing to LibTorch (an unpacked LibTorch download, or the one in the torch Python package: python3 -c "import torch; print(torch.utils.cmake_prefix_path)"). It supports CPU and CUDA builds; use a CUDA LibTorch with -DAMReX_GPU_BACKEND=CUDA. Build only this topic with -DAMReXTutorials_TOPICS=ML; see Getting Started.

PyTorchInference

Level:

★★★

Source:

ML/PyTorchInference

Key APIs:

torch::jit::load, torch::from_blob, c10::InferenceMode, c10::cuda::CUDAStreamGuard, Gpu::DeviceVector, Box::index, MFIter

A PyTorch model replaces a pointwise kernel: for every cell it maps two input fields to two output fields. The cells of each box are gathered into a batch, wrapped in a tensor without a copy, passed through the model, and the result is scattered back into a MultiFab. In CUDA builds, the data stay on the GPU and torch runs on AMReX’s CUDA stream.

make_model.py writes model.pt, a TorchScript model with one dense layer, \(y = \tanh(W x + b)\), with fixed weights, so that the program can compare the model output with this formula. A trained network is saved and used the same way, as long as it maps a (cells, inputs) tensor to a (cells, outputs) tensor.

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(2, 2)
        with torch.no_grad():
            # Must match the weights in main.cpp
            self.linear.weight.copy_(torch.tensor([[1.0, -0.5], [0.25, 2.0]]))
            self.linear.bias.copy_(torch.tensor([0.5, -0.25]))

    def forward(self, x):
        # x has shape (number of cells, 2)
        return torch.tanh(self.linear(x))

The model is loaded on the device where the MultiFab data live (the CPU, or the current CUDA device), and converted to the tensor type of amrex::Real:

// Load the TorchScript model written by make_model.py.
torch::jit::script::Module model;
try {
    model = torch::jit::load(model_file, device);
} catch (c10::Error const& e) {
    amrex::Abort("Failed to load " + model_file + ": " + e.what());
}
model.to(dtype);
model.eval();

For each box, the cells are gathered into an (ncells, nin) batch, the model is evaluated, and the output is scattered back:

// Gather the cells of the box into an (ncells, nin) batch in GPU
// memory for GPU builds.
Gpu::DeviceVector<Real> batch(ncells*nin);
Real* pbatch = batch.data();
ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
    auto const m = bx.index(IntVect(AMREX_D_DECL(i,j,k)));
    for (int n = 0; n < nin; ++n) {
        pbatch[m*nin+n] = in(i,j,k,n);
    }
});

// Wrap the batch in a tensor without copying, and evaluate the model.
torch::Tensor input = torch::from_blob(pbatch, {ncells, nin}, options);
torch::Tensor output;
{
    BL_PROFILE("model.forward");
    output = model.forward({input}).toTensor().contiguous();
}
AMREX_ALWAYS_ASSERT(output.dim() == 2 && output.size(0) == ncells &&
                    output.size(1) == nout && output.scalar_type() == dtype);

// Scatter the result back into the MultiFab.
Real const* pout = output.const_data_ptr<Real>();
ParallelFor(bx, [=] AMREX_GPU_DEVICE (int i, int j, int k)
{
    auto const m = bx.index(IntVect(AMREX_D_DECL(i,j,k)));
    for (int n = 0; n < nout; ++n) {
        out(i,j,k,n) = pout[m*nout+n];
    }
});

Configure from the repository root, build and run (see Getting Started):

cmake -S . -B build -DAMReXTutorials_TOPICS=ML -DAMReXTutorials_TORCH=ON \
      -DCMAKE_PREFIX_PATH=/path/to/libtorch
cmake --build build --target ML_PyTorchInference -j 8
cd build/ML/PyTorchInference
mpiexec -n 2 ./PyTorchInference inputs

model.pt is in the repository; regenerate it with python3 make_model.py (needs the torch Python package). With GNU Make, pass PYTORCH_ROOT=/path/to/libtorch.

Loaded model.pt on cpu
Output range: [-0.7585439934, 0.9635137063] and [-0.9863230868, 0.9632540745]
Max error vs. tanh(W x + b): 1.110223025e-16
PASSED