CosmoFlow: Using Deep Learning to Learn the Universe at Scale

Abstract

Deep learning is a promising tool to determine the physical model that describes our universe. To handle the considerable computational cost of this problem, we present CosmoFlow: a highly scalable deep learning application built on top of the TensorFlow framework. CosmoFlow uses efficient implementations of 3D convolution and pooling primitives, together with improvements in threading for many element-wise operations, to improve training performance on Intel(C) Xeon Phi(TM) processors. We also utilize the Cray PE Machine Learning Plugin for efficient scaling to multiple nodes. We demonstrate fully synchronous data-parallel training on 8192 nodes of Cori with 77% parallel efficiency, achieving 3.5 Pflop/s sustained performance. To our knowledge, this is the first large-scale science application of the TensorFlow framework at supercomputer scale with fully-synchronous training. These enhancements enable us to process large 3D dark matter distribution and predict the cosmological parameters ${\Omega }_M$, ${\sigma }_8$ and n${}_s$ with unprecedented accuracy.

11 pages, 6 pages, presented at SuperComputing 2018

References in corpus (6)

• HOGWILD!: A Lock-Free Approach to Parallelizing Stochastic Gradient Descent
• cuDNN: Efficient Primitives for Deep Learning
• Rotation-invariant convolutional neural networks for galaxy morphology prediction
• Large Batch Training of Convolutional Networks
• Solving Large Scale Structure in Ten Easy Steps with COLA
• Scaling GRPC Tensorflow on 512 nodes of Cori Supercomputer

Cited by in corpus (11)

• Learning to Predict the Cosmological Structure Formation
• A deep learning approach to cosmological dark energy models
• Baryon acoustic oscillations reconstruction using convolutional neural networks
• Adaptive Neural Network-Based Approximation to Accelerate Eulerian Fluid Simulation
• The bias of dark matter tracers: assessing the accuracy of mapping techniques
• One simulation to have them all: performance of the Bias Assignment Method against N-body simulations
• Using Machine Learning at Scale in HPC Simulations with SmartSim: An Application to Ocean Climate Modeling
• Sentinel: Runtime Data Management on Heterogeneous Main MemorySystems for Deep Learning
• Predicting Localized Primordial Star Formation with Deep Convolutional Neural Networks
• Derivation and Analysis of Fast Bilinear Algorithms for Convolution
• Cosmological parameter estimation from large-scale structure deep learning