Skip to content
This repository has been archived by the owner on Mar 10, 2020. It is now read-only.
/ caffe-admm Public archive

DNN connection pruning with Alternating Direction Method of Multipliers (ADMM)

License

Notifications You must be signed in to change notification settings

kaiqzhan/caffe-admm

Repository files navigation

Installation

Same as the original version of Caffe.

See http://caffe.berkeleyvision.org/installation.html for the latest installation instructions.

How to train

Two steps are required to prune a neural network. Suppose we have a pretrained model "bvlc_alexnet.caffemodel".

Step 1: admm (suppress target weights)

./build/tools/caffe train --solver=models/bvlc_alexnet/solver_admm.prototxt --weights=/path/to/bvlc_alexnet.caffemodel

Step 2: retrain (set target weights to 0 and retrain)

./build/tools/caffe train --solver=models/bvlc_alexnet/solver_retrain.prototxt --weights=models/bvlc_alexnet/caffe_alexnet_train_admm_iter_2400000.caffemodel

Suppose caffe_alexnet_train_admm_iter_2400000.caffemodel is the trained caffemodel from step 1.

Explanation of parameters

solver prototxt

  • lr_policy: "admm" for admm step; your prefered policy for retrain step
  • gamma: drop the learning rate by a factor of this value
  • stepvalue: drop the learning rate by gamma when meet this stepvalue in each admm iteration
  • admm_iter: Caffe training iterations for each admm iteration
  • max_iter: Total training iterations (total admm iterations = max_iter / admm_iter)
  • pruning_phase: "admm" for admm step and "retrain" for retrain step

model prototxt

  • prune_ratio: the percent of weight parameters to prune for each layer
  • rho: the rho parameter for admm pruning (detailed in our paper)

Caffe

Build Status License

Caffe is a deep learning framework made with expression, speed, and modularity in mind. It is developed by Berkeley AI Research (BAIR)/The Berkeley Vision and Learning Center (BVLC) and community contributors.

Check out the project site for all the details like

and step-by-step examples.

Custom distributions

Community

Join the chat at https://gitter.im/BVLC/caffe

Please join the caffe-users group or gitter chat to ask questions and talk about methods and models. Framework development discussions and thorough bug reports are collected on Issues.

Happy brewing!

License and Citation

Caffe is released under the BSD 2-Clause license. The BAIR/BVLC reference models are released for unrestricted use.

Please cite Caffe in your publications if it helps your research:

@article{jia2014caffe,
  Author = {Jia, Yangqing and Shelhamer, Evan and Donahue, Jeff and Karayev, Sergey and Long, Jonathan and Girshick, Ross and Guadarrama, Sergio and Darrell, Trevor},
  Journal = {arXiv preprint arXiv:1408.5093},
  Title = {Caffe: Convolutional Architecture for Fast Feature Embedding},
  Year = {2014}
}

About

DNN connection pruning with Alternating Direction Method of Multipliers (ADMM)

Resources

License

Stars

Watchers

Forks

Packages

No packages published