兆体素域上的分布式多重网格神经求解器
机器学习
2021-04-30 v1 机器学习
摘要
我们考虑用于作为 PDE 求解器并产生全场输出的大规模神经网络的分布式训练。具体而言，我们考虑兆体素域上广义三维泊松方程（3D Poisson equation）的神经求解器。我们提出一个可扩展框架，集成了两项不同进展。首先，我们通过类比数值线性代数中使用的多重网格（multigrid）技术来加速大模型训练。此处，网络依次使用递增分辨率的输入层次进行训练，类比多重网格方法中使用的 'V'、'W'、'F' 和 'Half-V' 循环。结合多重网格方法，我们实现了分布式深度学习框架，显著减少求解时间。我们在 GPU（云上 Azure VM）与 CPU 集群（PSC Bridges2）上展示了该方法的可扩展性。该方案被部署用于训练一个广义三维泊松求解器，可良好扩展至为高位维输入族预测分辨率达 512x512x512 的输出全场解。
引用
@article{arxiv.2104.14538,
title = {Distributed Multigrid Neural Solvers on Megavoxel Domains},
author = {Aditya Balu and Sergio Botelho and Biswajit Khara and Vinay Rao and Chinmay Hegde and Soumik Sarkar and Santi Adavani and Adarsh Krishnamurthy and Baskar Ganapathysubramanian},
journal= {arXiv preprint arXiv:2104.14538},
year = {2021}
}