alexfjw / prunnable-layers-pytorch

Prunable nn layers for pytorch.

Prunable neural network layers for Pytorch

The code in this repo has been extracted from my ML udacity capstone project here, where I experimented with CNNs for Japanese OCR.

Prunable versions of common CNN layers can be found in this repo. We use taylor expansion to estimate the importance of feature maps, as described by Molchanov, P. et al[1].
A good summary of this approach can be found here.

Pruning & taylor estimation was written with an OOP approach for readability.

Dependencies

• Python 3.6.3
• Pytorch 0.4
• torchvision (for vgg example)

Description of code

The layers in prunable_nn.py are plug and play. However, they only compute the importance of each feature map. (pruning is a concept linked to models, and not individual layers). Filtering to select the smallest feature map and dropping inputs for layer next to the pruned layer has to be done manually.

PConv2d is a Conv2d layer which registers a backward hook during gradient calulation to weigh the importance of each feature map.

Feature maps are stored after every forward call, and talyor estimates for feature map importance is calculated after every backward call.

As mentioned, removing a feature map reduces the outputs of the layer. The next immediate layer has to take in fewer inputs. PLinear and PBatchNorm have been written to perform this.

Usage Examples

Simply use PLinear(..), PConv2(..) & PBatchNorm(...) in place of the originals.
Also modify your models to add a prune method.
See models.py for further information.

The tests in ./tests/ may be useful for testing your prunable models.

For benchmarks & examples of the full training & pruning procedure, refer to my ML udacity capstone project here.

Possible Extensions

Making PLinear prunable - shouldn't be too difficult, implementation should be identical to PConv2d.

References

[1] Molchanov, P., Tyree, S., Karras, T., Aila, T., & Kautz, J. (2017, June 08). Pruning Convolutional Neural Networks for Resource Efficient Inference. Retrieved December 10, 2017, from https://arxiv.org/abs/1611.06440

prunnable-layers-pytorch's Issues

Shouldn't `taylor_estimates` calculate aggregated ranking instead of ranking of last forward pass?

It seems that in this line, taylor_estimates is overwritten after every pass. But I believe the channel ranking should be based on not just one pass but the average of multiple passes. How is this managed in the current implementation? Is it the case that the logic to calculate aggregate ranking can be implemented outside this model?

Thanks for the code by the way. It's pretty helpful.