A Deep Learning Based Physics Informed Continuous Spatio Temporal Super-Resolution Framework

American Physical Society Conference, Nov. 2020

Authors - Soheil Esmaeilzadeh*, Chiyu “Max” Jiang*, Kamyar Azizzadenesheli, Karthik Kashinath, Mustafa Mustafa, Hamdi A. Tchelepi, Philip Marcus, Prabhat, and Anima Anandkumar

*denotes equal contribution

Abstract - We propose a novel deep learning based super-resolution framework to generate continuous (grid-free) spatio-temporal solutions from the low-resolution inputs. While being computationally efficient, our proposed framework accurately recovers the fine-scale quantities of interest and allows for: (i) the output to be sampled at all spatio-temporal resolutions, (ii) a set of Partial Differential Equation (PDE) constraints to be imposed, and (iii) training on fixed-size inputs on arbitrarily sized spatio-temporal domains owing to its fully convolutional encoder. We empirically study the performance of our framework on the task of super-resolution of turbulent flows in the Rayleigh-Benard convection problem. Across a diverse set of evaluation metrics, we show that our proposed framework significantly outperforms the existing baselines. Furthermore, we provide a large-scale implementation of our framework and show that it efficiently scales across large clusters, achieving 96.80 percent scaling efficiency on up to 128 GPUs and a training time of less than 4 minutes. We provide an open-source implementation of our method that supports arbitrary combinations of PDE constraints.

Keywords: Super-Resolution, PDEs, Physics-Constrained, Deep Neural Networks

Sample tuples of the ground truth high-resolution data, the low-resolution input data for MeshfreeFlowNet, and the high-resolution super-resolved data by MeshfreeFlowNet for the 4 physical parameters of the Rayleigh–Bénard problem namely $T$, $p$, $u$, $w$ as the temperature, pressure, and the x and z components of the velocity.

Highlights

Fig.1. Schematic for the training pipeline for MeshfreeFlowNet model for continuous space-time super-resolution. An input low-resolution grid is fed to the Context Generation Network that creates a Latent Context Grid. A random set of points in the corresponding space-time domain is sampled to query the Latent Context Grid, and the physical output values at these query locations can be continuously decoded using a Continuous Decoding Network, implemented as a Multilayer Perception. Due to the differentiable nature of the MLP, any partial derivatives of the output physical quantities with respect to the input space-time coordinates can be effectively computed via backpropagation, that can be combined with the PDE constraints to produce an Equation Loss. On the other hand, the predicted value can be contrasted with the ground truth value at these locations produced by interpolating the high-resolution ground truth to produce a Prediction Loss. Gradients from the combined losses can be backpropagated to the network for training.

Fig.2. A schematic for the architecture of the MeshfreeFlowNet framework. MeshfreeFlowNet consists of two end-to-end trainable modules: the Context Generation Network, and the Continuous Decoding Network. The Baseline (II) method that we exhaustively compare with share the same 3D U-Net structure in the encoding arm, but instead uses a convolutional decoder for producing the final output. In comparison, the MeshfreeFlowNet allows for continuous output that can be sampled at arbitrary space-time locations.

Fig.3. Comparison between the performance of MeshfreeFlowNet framework for super-resolving the low-resolution data vs. two baseline models.

Fig.4. For a MeshfreeFlowNet model that has been respectively trained on 1 dataset and 10 datasets each having a different initial condition, the super-resolution performance evaluation is reported on a dataset with an unseen initial condition.

Fig.5. For a MeshfreeFlowNet model that has been trained on 10 datasets each having a different boundary condition (Rayleigh number) as Ra ∈ [2, 90] × 105 with Pr = 1, the super-resolution performance evaluation is reported for: a Rayleigh number within the range of boundary conditions of the training sets (i.e., Ra = 5×106), Rayleigh numbers slightly below and above the range of boundary conditions of the training sets (i.e., Ra = 1 × 105 and Ra = 1 × 107 respectively), and Rayleigh numbers far below and above the range of boundary conditions of the training sets (i.e., Ra = 1 × 104 and Ra = 1 × 108 respectively).

Fig.6. Scaling performance on up to 128 GPUs on NERSC Cori GPU cluster. Scaling efficiency is close to ideal performance, at 96% with all 128 GPUs. Training convergence is greatly accelerated by scaling across more GPUs. The shaded band in Figs. (b)-(c) denotes the (10%, 90%) noise interval. Throughput refers to the number of samples per second.