musco-ai/musco-tf

MUSCO: Multi-Stage COmpression of neural networks

MUSCO: Multi-Stage COmpression of neural networks

This repository contains supplementary code for the paper MUSCO: Multi-Stage COmpression of neural networks. It demonstrates how a neural network with convolutional and fully connected layers can be compressed using iterative tensor decomposition of weight tensors.

Requirements

numpy
scipy
scikit-tensor-py3
tensorly-musco
absl-py
tqdm
tensorflow-gpu (TensorRT support)

Installation

pip install musco-tf

Quick Start

from musco.tf import CompressorVBMF, Optimizer

model = load_model("model.h5")
compressor = CompressorVBMF(model)

while True:
    model = compressor.compress_iteration(number=5)
    
    # Fine-tune compressed model.

# Compressor decomposes 5 layers on each iteration
# and returns compressed model. You have to fine-tune
# model after each iteration to restore accuracy.
# Compressor automatically selects the best parameters
# for decomposition on each iteration.

# You can freeze and quantize model after compression.
optimizer = Optimizer(precision="FP16", max_batch_size=16)
optimizer.freeze(model)
optimizer.optimize("frozen.pb")

Citing

If you used our research, we kindly ask you to cite the corresponding paper.

@article{gusak2019one,
  title={MUSCO: Multi-Stage Compression of neural networks},
  author={Gusak, Julia and Kholiavchenko, Maksym and Ponomarev, Evgeny and Markeeva, Larisa and Oseledets, Ivan and Cichocki, Andrzej},
  journal={arXiv preprint arXiv:1903.09973},
  year={2019}
}

License

Project is distributed under Apache License 2.0.

musco-tf's Issues

Upgrade to tf 2.8.2

For people who are trying to use this package with a new version of tensorflow. These are some of the changements that should be done in order to run this code using latest numpy and tf versions. If you have any other changements please feel free to add them.

Steps:

• Clone scikit-tensor-py3 Git repo and delete all np.float statements. Replace them with float().
• Install the repo again locally by pip install .
• Replace
  • from tensorflow.python.keras.engine import InputSpec
  • by
  • from tensorflow.keras.layers import InputSpec
• Replace trt usage
  • import tensorflow.contrib.tensorrt as trt
  • by
  • from tensorflow.python.compiler.tensorrt import trt_convert as trt
• Delete all tf.Session and tf.reset_default_graph()

Installation fails due to conflicting numpy version

Hi, users are unable to run musco-tf due to dependency conflict with numpy package. As shown in the following full dependency graph of musco-tf, musco-tf requires numpy *，while scikit-tensor-py3 requires numpy ==1.16.*.

According to pip’s “first found wins” installation strategy, numpy 1.17.4 is the actually installed version. However, numpy 1.17.4 does not satisfy ==1.16.*.

Dependency tree-----------

musco-tf - 1.0.2
| +- absl-py(install version:0.8.1 version range:*)
| +- numpy(install version: 1.17.4 version range:*)
| +- scikit-tensor-py3(install version:0.4.1 version range:*)
| | +- numpy(install version:1.16.5 version range:==1.16.*)
| | +- scipy(install version:1.3.3 version range:==1.3.)
| +- scipy(install version:1.3.3 version range:*)
| +- tensorly-musco(install version:0.4.5 version range:*)
| +- tqdm(install version:4.40.2 version range:*)

Thanks for your help.
Best,
Neolith