Distributed Multigrid Neural Solvers on Megavoxel Domains

Thumbnail Image
Date
2021-01-01
Authors
Balu, Aditya
Botelho, Sergio
Khara, Biswajit
Rao, Vinay
Sarkar, Soumik
Hegde, Chinmay
Adavani, Santi
Major Professor
Advisor
Committee Member
Journal Title
Journal ISSN
Volume Title
Publisher
Authors
Person
Krishnamurthy, Adarsh
Associate Professor
Person
Research Projects
Organizational Units
Organizational Unit
Mechanical Engineering
The Department of Mechanical Engineering at Iowa State University is where innovation thrives and the impossible is made possible. This is where your passion for problem-solving and hands-on learning can make a real difference in our world. Whether you’re helping improve the environment, creating safer automobiles, or advancing medical technologies, and athletic performance, the Department of Mechanical Engineering gives you the tools and talent to blaze your own trail to an amazing career.
Organizational Unit
Journal Issue
Is Version Of
Versions
Series
Department
Mechanical EngineeringElectrical and Computer EngineeringPlant Sciences Institute
Abstract

We consider the distributed training of large-scale neural networks that serve as PDE solvers producing full field outputs. We specifically consider neural solvers for the generalized 3D Poisson equation over megavoxel domains. A scalable framework is presented that integrates two distinct advances. First, we accelerate training a large model via a method analogous to the multigrid technique used in numerical linear algebra. Here, the network is trained using a hierarchy of increasing resolution inputs in sequence, analogous to the 'V', 'W', 'F', and 'Half-V' cycles used in multigrid approaches. In conjunction with the multi-grid approach, we implement a distributed deep learning framework which significantly reduces the time to solve. We show the scalability of this approach on both GPU (Azure VMs on Cloud) and CPU clusters (PSC Bridges2). This approach is deployed to train a generalized 3D Poisson solver that scales well to predict output full-field solutions up to the resolution of 512x512x512 for a high dimensional family of inputs.

Comments

This is a pre-print of the article Balu, Aditya, Sergio Botelho, Biswajit Khara, Vinay Rao, Chinmay Hegde, Soumik Sarkar, Santi Adavani, Adarsh Krishnamurthy, and Baskar Ganapathysubramanian. "Distributed Multigrid Neural Solvers on Megavoxel Domains." arXiv preprint arXiv:2104.14538 (2021). Posted with permission.

Description
Keywords
Citation
DOI
Source
Copyright
Fri Jan 01 00:00:00 UTC 2021
Collections