Skip to content

Pytorch version for weight pruning for Murata Group's CREST project

Notifications You must be signed in to change notification settings

zepx/pytorch-weight-prune

Repository files navigation

Repository Features

The goal of this repository is to apply Pruning for AlexNet and potentially, future models. The original implementation does not work with multiple-gpu, which has been fixed in this repository.

Neural Network Pruning PyTorch Implementation

Original Authors: Luyu Wang & Gavin Ding

Original Repository: https://github.com/wanglouis49/pytorch-weights_pruning.

Borealis AI

Motivation

Neural network pruning has become a trendy research topic, but we haven't found an easy to use PyTorch implementation. We want to take advantage of the power of PyTorch and build pruned networks to study their properties.

Note: this implementation is not aiming at obtaining computational efficiency but to offer convenience for studying properties of pruned networks. Discussions on how to have an efficient implementation is welcome. Thanks!

High-level idea

  1. We write wrappers on PyTorch Linear and Conv2d layers.
  2. For each layer, once a binary mask tensor is computed, it is multiplied with the actual weights tensor on the forward pass.
  3. Multiplying the mask is a differentiable operation and the backward pass is handed by automatic differentiation (no explicit coding needed).

Pruning methods

Weight pruning

Han et al propose to compress deep learning models via weights pruning Han et al, NIPS 2015. This repo is an implementation in PyTorch. The pruning method is replaced by the "class-blinded" method mentioned in See et al, CoNLL 2016, which is much easier to implement and has better performance as well.

Filter pruning

Pruning convolution filters has the advantage that it is more hardware friendly. We also implement the "minimum weight'' approach in Molchanov et al, ICLR 2017

About

Pytorch version for weight pruning for Murata Group's CREST project

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages