wxb506 / sparse-winograd-cnn

Efficient Sparse-Winograd Convolutional Neural Networks

Efficient Sparse-Winograd Convolutional Neural Networks

This is the code and models for paper Efficient Sparse-Winograd Convolutional Neural Networks by Xingyu Liu et al.

Introduction

This work is based on our ICLR 2018 paper. We propose modifications to Winograd-based CNN architecture to enable operation savings from Winograd’s minimal filtering algorithm and network pruning to be combined.

Convolutional Neural Networks (CNNs) are computationally intensive, which limits their application on mobile devices. Their energy is dominated by the number of multiplies needed to perform the convolutions. Winograd’s minimal filtering algorithm and network pruning can reduce the operation count, but these two methods cannot be straightforwardly combined — applying the Winograd transform fills in the sparsity in both the weights and the activations. We propose two modifications to Winograd-based CNNs to enable these methods to exploit sparsity.

In this repository, we release code and data for training Winograd-ReLU CNN on ImageNet as well as pre-trained and iteratively pruned Winograd-ReLU models.

Citation

If you find our work useful in your research, please consider citing:

@article{Liu2018Winograd,
  title={Efficient Sparse-Winograd Convolutional Neural Networks},
  author={Xingyu Liu and Jeff Pool and Song Han and William J. Dally},
  journal={International Conference on Learning Representations (ICLR)},
  year={2018}
}

Installation

Install TensorFlow and Tensorpack. The code has been tested with Python 2.7, TensorFlow 1.3.0, CUDA 8.0 and cuDNN 5.1 on Ubuntu 14.04.

Users may also need to download raw ImageNet dataset for ImageNet experiments. Users should ensure that the Tensorpack ResNet example can run with ImageNet.

Install customized Tensorflow Op:

cd /path/to/Sparse-Winograd-CNN/winograd2x2_cublas
make
export PYTHONPATH=/path/to/Sparse-Winograd-CNN/winograd2x2_cublas:$PYTHONPATH

Users may also change the -arch flag in winograd2x2_cublas/winograd2x2_imTrans/Makefile and winograd2x2_cublas/winograd2x2_conv/Makefile to suit their GPU computing capability.

Put ResNet-18-var/winograd_conv.py and ResNet-18-var/winograd_imtrans.py into the cloned tensorpack/models directory.

Usage

To train the Winograd-ReLU CNN from scratch on ImageNet with GPU 0 and 1:

./imagenet-resnet-transWino-prune.py --gpu 0,1 --data /path/to/dataset/imagenet

To use pre-trained model or test with pruned model, download the models. Then run with command:

./imagenet-resnet-transWino-prune.py --gpu 0,1 --data /path/to/dataset/imagenet --load /path/to/model-name.data-00000-of-00001

We also provided scripts for pruning, retraining and viewing the model: ResNet-18-var/prune_sh.sh, retrain_sh.sh and view_sh.sh.

License

Our code is released under MIT License (see LICENSE file for details).