aygong / snip-tensorflow

Code for the project "SNIP: Single-Shot Network Pruning"

CS-439 Optimization for Machine Learning - Class Project

    

Group Members

• Aoyu Gong ([email protected])
• Sijia Du ([email protected])
• Qiyuan Dong ([email protected])

Introduction

This project focuses on a network pruning algorithm proposed on ICLR 2019, namely SNIP for single-shot network pruning. We reproduced the experiments of the original paper with a similar experimental setup. In particular, we evaluated the algorithm on two datasets with seven network architectures given by the reference. Our implementation is based on the GitHub repo provided by the authors of the original paper. To further study its robustness, we validated the effectiveness of the algorithm on an additional dataset and nine network architectures, provided numerical and visualization results with respect to a wide range of configurations, and compared our result with the original paper.

Requirements

Dependencies

You can reproduce our results by two ways, using either Google Colab Pro or your local system.

• Reproduce our results on Colab.

  • Colab contains most of needed packages, and is easy and highly recommended for reproducing our results.

  • Note that our codes requires tensorflow-gpu 1.14.0, which is not pre-installed in Colab. Be sure to run

    # Install required packages
    %tensorflow_version 2.x
    !pip install gast==0.2.2
    !pip uninstall -y tensorflow
    !pip install tensorflow-gpu==1.14.0

    at the beginning of the Colab notebook run.ipynb to setup the environment correctly.

• Reproduce our results on your local system.

  • To reproduce our results, you should install the following packages.
    • gast 0.2.2
    • matplotlib 3.3.4
    • numpy 1.19.2
    • imgaug 0.4.0
    • tensorflow-gpu 1.14.0
  • GPUs with at least 25.46 GB RAM are required. Otherwise, you will get RuntimeError: CUDA out of memory.

Datasets

Download the following datasets and move the files to the corresponding folders.

• Download the Kuzushiji-MNIST dataset and move the files into ./KMNIST.
• Download the MNIST dataset and move the files into ./MNIST.
• Download the CIFAR-10 dataset and move the files into ./cifar-10-batches-py.

How to Reproduce Our Experiments

You can reproduce our experiments in two ways, using either Google Colab Pro or your local machine.

• Reproduce our experiments on Colab.

  • First, create a folder called Mini-Project in your Google Drive, i.e., /gdrive/MyDrive/Mini-Project.

  • Second, make sure the structure of the folder is the same as the Folder Structure section.

    • git clone our repo to your local machine and then move the code to the folder.
    • Download the three datasets and then move the files to the folder.

  • Third, click Change runtime type under the menu Runtime, and select GPU as Hardware accelerator and High-RAM as Runtime shape.

  • Last, run the provided Jupyter notebook run.ipynb under the root directory with Run all.

• Reproduce our experiments on your local machine.

  • First, install all the dependencies mentioned in the Requirement section.
  • Second, create a folder called Mini-Project on your local machine, i.e., ./Mini-Project.
  • Third, make sure the structure of the folder is the same as the Folder Structure section.
    • git clone our repo to your local machine and then move the code to the folder.
    • Download the three datasets and then move the files to the folder.
  • Last, run the provided Jupyter notebook run.ipynb under the root directory with Run all.

Folder Structure

./Mini-Project
├── README.md
|
├── results.xlsx                    # All the results from our experiments
|
├── report.pdf                      # Our report
|
├── run.ipynb                       # The Jupyter notebook to run the algorithm and reproduce our results.
|
├── snip/                           # Folder contains all the sources files to run the algorithm
|   ├── train.py                    # Define the training process
|   ├── test.py                     # Define the test process
|   ├── prune.py                    # Define the pruning process which is executed before training
|   ├── network.py                  # Define multiple neural network architectures
|   ├── model.py                    # Define the entire model, including the network, loss, optimization, and input/output processing, etc.
|   ├── mnist.py                    # Process the MNIST dataset
|   ├── main.py                     # Main file that parses arguments and initializes the entire training and testing procedure
|   ├── kmnist.py                   # Process the KMNIST dataset
|   ├── helpers.py                  # Several auxiliary functions
|   ├── dataset.py                  # Process the dataset to feed the model
|   ├── cifar.py                    # Process the CIFAR-10 dataset
|   └── augment.py                  # Perform data augmentation
|
├── cifar-10-batches-py/            # Folder for the CIFAR-10 dataset
|   ├── batches.meta
|   ├── data_batch_1
|   ├── data_batch_2
|   ├── data_batch_3
|   ├── data_batch_4
|   ├── data_batch_5
|   ├── readme.html
|   └── test_batch
|
├── KMNIST/                         # Folder for the KMNIST dataset
|   ├── t10k-images-idx3-ubyte.gz
|   ├── t10k-labels-idx1-ubyte.gz
|   ├── train-images-idx3-ubyte.gz
|   └── train-labels-idx1-ubyte.gz
|
└── MNIST/                          # Folder for the MNIST dataset
    ├── t10k-images-idx3-ubyte.gz
    ├── t10k-labels-idx1-ubyte.gz
    ├── train-images-idx3-ubyte.gz
    └── train-labels-idx1-ubyte.gz

Useful Links

The algorithm of SNIP is presented in our report and the original paper.

The following network architectures are used for evaluating.

• resnet-18, resnet-34: Deep Residual Learning for Image Recognition
• squeezenet-vanilla, squeezenet-bypass: SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and < 0.5 MB model size
• googlenet: Going Deeper with Convolutions
• densenet-121, densenet-169, densenet-201, densenet-264: Densely Connected Convolutional Networks