Libtorch (PyTorch C++ API)

MOOSE can be linked against libtorch Paszke et al. (2019) to enable hardware acceleration and some high-level machine-learning capabilities. The libtorch source is provided as a git submodule under framework/contrib/pytorch, and a setup script is provided to build and install it with the configuration MOOSE expects.

Prerequisites

The libtorch build depends on PETSc so that BLAS and LAPACK are consistent between the two libraries. PETSc must therefore be installed before libtorch. The script looks for PETSc at ${PETSC_DIR} (defined if you are using conda, and otherwise defaults to <MOOSE_DIR>/petsc/arch-moose) and aborts if it cannot be found.

A working compiler stack and CMake are also required. The MOOSE conda environment satisfies both.

Installation

To build and install libtorch, run


cd ~/projects/moose
./scripts/update_and_rebuild_libtorch.sh
schooltip

The setup script uses sensible defaults that work out-of-the-box. The script is also extensively customizable. Use the --help argument to print out a detailed help message.

The most commonly used options and environment variables are:

  • --fast — skip the update, clean, and configure steps (re-build only).

  • --skip-submodule-update — skip the submodule update step.

  • LIBTORCH_DIR — installation directory. Defaults to <MOOSE_DIR>/framework/contrib/pytorch/installed.

  • LIBTORCH_SRC_DIR — use a custom libtorch source tree instead of the bundled submodule. Setting this implies --skip-submodule-update.

  • LIBTORCH_JOBS — number of parallel build jobs. Defaults to MOOSE_JOBS, or 1 if unset.

  • PETSC_DIR — path to the PETSc installation.

commentnote:GPU support

CUDA is enabled automatically when a CUDA toolkit is detected on the system. For Intel GPU (XPU) support, export USE_XPU=1 before invoking the script. General CMake environment variables are also respected, so other backends can be toggled by passing additional -D... arguments through to the configure step.

Building libtorch is resource-intensive: expect a multi-gigabyte build directory and a long compile time on the first build. Subsequent rebuilds can be sped up with --fast.

Configuring MOOSE with libtorch

Once libtorch has been installed, configure MOOSE with


./configure --with-libtorch

If libtorch was installed to a non-default location, pass that path explicitly:


./configure --with-libtorch=/path/to/libtorch

HPC systems

On INL HPC systems, libtorch is already provided through the moose-dev module/container and does not need to be built manually. See the INL HPC installation instructions for more information. On other HPC systems, follow the same procedure as for a workstation, taking care to load a compatible compiler stack and CMake before invoking the script.

References

  1. Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: an imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d\textquotesingle Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 8024–8035. Curran Associates, Inc., 2019. URL: http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf.[BibTeX]