TorchBraid
This package implements a layer-parallel approach to training neural ODEs, and neural networks
Area: Artificial intelligence
CASS member: LEADS
Description
Algorithmically multigrid-in-time is used to expose parallelism in the forward and backward propagation phases used to compute the gradient. The neural network interface is build on PyTorch, while the backend uses XBraid (a C library) for multigrid-in-time. TorchBraid uses direct GPU communication when running simulations on GPUs, a CUDA-aware MPI implementation is required for multi-GPU support.
Target audience
Science teams interested layer-parallel approach to training of neural networks.
License: BSD-3-Clause
Package links
- E4S: TORCHBRAID