Distributed Fourier Neural Operators

Fourier Neural Operators (FNOs) are a class of neural operator, which use weightings on the Fourier transform of their inputs to approximate infinite-dimensional mappings between function spaces. They are particularly useful in approximating the solutions of smooth parametric (e.g. by permeability) partial differential equations (PDEs), and once trained are capable of producing a nearly identical output to a traditional numerical solver roughly three orders of magnitude faster, making them very useful in a wide variety of engineering applications that require repeated PDE solves (e.g. during uncertainty quantification). Until now, FNOs have been limited to small problems, as their memory intensive design makes them difficult to scale in a traditional machine learning setting on a single computer or GPU. In this work, a decomposition scheme is described and implemented in PyTorch using the DistDL distributed deep learning framework, which is capable of scaling FNOs to arbitrary dimension and input size running on many nodes in a distributed memory system. This parallel implementations allows FNOs to be trained and run on problems of a practical scale on both CPU and GPU clusters.