Roulette: A pruning framework to train a sparse neural network from scratch

Abstract

Due to space and inference time restrictions, finding an efficient and sparse sub-network from a dense and over-parameterized network is critical for deploying neural networks on edge devices. Recent efforts explore obtaining a sparse sub-network by performing network pruning during training procedures to reduce training costs, such as memory and floating-point operations (FLOPs). However, these works take more than $1.4\times $ the total number of iterations and try all possible pruning parameters manually to obtain sparse sub-networks. In this paper, we present a pruning framework Roulette to train a sparse network from scratch. First, we propose a novel method to train a sparse network by Pruning through the lens of Loss Landscape iteratively and automatically (PLL). We do a theoretical analysis that the curvature of the loss function is higher in the initial phase and can conduct us to start network pruning. According to our results on CIFAR-10/100 and ImageNet dataset, PLL saves up to $4\times $ training FLOPs than prior works while maintaining comparable or even better accuracy. Then we design push and pull operations to synchronize the pruned weights on different GPUs during training, scaling PLL to multiple GPUs linearly. To our knowledge, Roulette is the first network pruning framework supporting multiple GPUs linearly.