Only Train Once: A One-Shot Neural Network Training And Pruning Framework

TL;DR

This paper introduces a novel one-shot training and pruning framework for deep neural networks that enables simultaneous training and compression without fine-tuning, achieving state-of-the-art results on multiple benchmarks.

Contribution

The proposed OTO framework combines zero-invariant grouping and a new optimization algorithm to prune networks during training, eliminating the need for fine-tuning after pruning.

Findings

Achieves significant FLOPs reduction and parameter compression.

State-of-the-art results on CIFAR10, SQuAD, and competitive on ImageNet.

Outperforms standard methods in group sparsity exploration.

Abstract

Structured pruning is a commonly used technique in deploying deep neural networks (DNNs) onto resource-constrained devices. However, the existing pruning methods are usually heuristic, task-specified, and require an extra fine-tuning procedure. To overcome these limitations, we propose a framework that compresses DNNs into slimmer architectures with competitive performances and significant FLOPs reductions by Only-Train-Once (OTO). OTO contains two keys: (i) we partition the parameters of DNNs into zero-invariant groups, enabling us to prune zero groups without affecting the output; and (ii) to promote zero groups, we then formulate a structured-sparsity optimization problem and propose a novel optimization algorithm, Half-Space Stochastic Projected Gradient (HSPG), to solve it, which outperforms the standard proximal methods on group sparsity exploration and maintains comparable convergence. To demonstrate the effectiveness of OTO, we train and compress full models simultaneously from scratch without fine-tuning for inference speedup and parameter reduction, and achieve state-of-the-art results on VGG16 for CIFAR10, ResNet50 for CIFAR10 and Bert for SQuAD and competitive result on ResNet50 for ImageNet. The source code is available at https://github.com/tianyic/only_train_once.