Distiller:神经网络压缩研究框架

Distiller是由Intel AI Lab维护的基于PyTorch的开源神经网络压缩框架。主要包括：

• 用于集成剪枝（pruning），正则化（regularization）和量化（quantization ）算法的框架。
• 一套用于分析和评估压缩性能的工具。
• 现有技术压缩算法的示例实现。

这算是目前我发现的最完整的压缩框架了，比较适合科研工作。下面简单说一下安装和使用。

安装

创建虚拟环境

Distiller是基于python的开源框架，为了不与其他工作冲突，最好提前创建一个虚拟环境。官方给的教程是使用virtualenv的，这个功能简单一些，这里有安装教程。目前流行的python包管理是conda（我用的也是），用的顺手就行。之后activate到创建的环境。

注意: – Distiller 仅在 Ubuntu 16.04 LTS中 Python 3.5环境下测试，并且默认使用了GPU，如果不使用GPU代码可能需要微小的调整。

克隆仓库并安装依赖

$ git clone https://github.com/NervanaSystems/distiller.git 

Clone Distiller的仓库。

$ pip3 install -r requirements.txt 

数据集

分类数据集的整理需要按照文档给定的方式（这里）。

distiller
  examples
    classifier_compression
data.imagenet/
    train/
    val/
data.cifar10/
    cifar-10-batches-py/
        batches.meta
        data_batch_1
        data_batch_2
        data_batch_3
        data_batch_4
        data_batch_5
        readme.html
        test_batch

当然还是比较常规的，建立软连接即可。

使用

仓库提供了大量的实例（distiller/examples/），本文以distiller/examples/classifier_compression/compress_classifier.py为例。

命令行参数

这个文件包含了一整套压缩框架，参数帮助可以调用以下命令：

$ python3 compress_classifier.py --help 

比如我们输入以下指令：

$ time python3 compress_classifier.py -a alexnet --lr 0.005 -p 50 ../../../data.imagenet -j 44 --epochs 90 --pretrained --compress=../sensitivity-pruning/alexnet.schedule_sensitivity.yaml Parameters: +----+---------------------------+------------------+---------------+----------------+------------+------------+----------+----------+----------+------------+---------+----------+------------+ | | Name | Shape | NNZ (dense) | NNZ (sparse) | Cols (%) | Rows (%) | Ch (%) | 2D (%) | 3D (%) | Fine (%) | Std | Mean | Abs-Mean | |----+---------------------------+------------------+---------------+----------------+------------+------------+----------+----------+----------+------------+---------+----------+------------| | 0 | features.module.0.weight | (64, 3, 11, 11) | 23232 | 13411 | 0.00000 | 0.00000 | 0.00000 | 0.00000 | 0.00000 | 42.27359 | 0.14391 | -0.00002 | 0.08805 | | 1 | features.module.3.weight | (192, 64, 5, 5) |  |  | 0.00000 | 0.00000 | 0.00000 | 1.91243 | 0.00000 | 62.38281 | 0.04703 | -0.00250 | 0.02289 | | 2 | features.module.6.weight | (384, 192, 3, 3) |  |  | 0.00000 | 0.00000 | 0.00000 | 6.18490 | 0.00000 | 61.33445 | 0.03354 | -0.00184 | 0.01803 | | 3 | features.module.8.weight | (256, 384, 3, 3) |  |  | 0.00000 | 0.00000 | 0.00000 | 6.96411 | 0.00000 | 64.38881 | 0.02646 | -0.00168 | 0.01422 | | 4 | features.module.10.weight | (256, 256, 3, 3) |  |  | 0.00000 | 0.00000 | 0.00000 | 15.49225 | 0.00000 | 68.30614 | 0.02714 | -0.00246 | 0.01409 | | 5 | classifier.1.weight | (4096, 9216) |  |  | 0.00000 | 0.21973 | 0.00000 | 0.21973 | 0.00000 | 90.99604 | 0.00589 | -0.00020 | 0.00168 | | 6 | classifier.4.weight | (4096, 4096) |  |  | 0.21973 | 3.46680 | 0.00000 | 3.46680 | 0.00000 | 89.37387 | 0.00849 | -0.00066 | 0.00263 | | 7 | classifier.6.weight | (1000, 4096) |  |  | 3.36914 | 0.00000 | 0.00000 | 0.00000 | 0.00000 | 75.71440 | 0.01718 | 0.00030 | 0.00778 | | 8 | Total sparsity: | - |  |  | 0.00000 | 0.00000 | 0.00000 | 0.00000 | 0.00000 | 88.43694 | 0.00000 | 0.00000 | 0.00000 | +----+---------------------------+------------------+---------------+----------------+------------+------------+----------+----------+----------+------------+---------+----------+------------+ 2018-04-04 21:30:52,499 - Total sparsity: 88.44 2018-04-04 21:30:52,499 - --- validate (epoch=89)----------- 2018-04-04 21:30:52,499 -  samples (256 per mini-batch) 2018-04-04 21:31:04,646 - Epoch: [89][ 50/ 500] Loss 2. Top1 51. Top5 74.023438 2018-04-04 21:31:06,427 - Epoch: [89][ 100/ 500] Loss 2. Top1 51. Top5 74. 2018-04-04 21:31:11,432 - Epoch: [89][ 150/ 500] Loss 2. Top1 51. Top5 74. 2018-04-04 21:31:14,364 - Epoch: [89][ 200/ 500] Loss 2. Top1 51. Top5 74. 2018-04-04 21:31:18,381 - Epoch: [89][ 250/ 500] Loss 2. Top1 51. Top5 74. 2018-04-04 21:31:22,195 - Epoch: [89][ 300/ 500] Loss 2. Top1 51. Top5 74. 2018-04-04 21:31:25,508 - Epoch: [89][ 350/ 500] Loss 2. Top1 51. Top5 74. 2018-04-04 21:31:29,538 - Epoch: [89][ 400/ 500] Loss 2. Top1 51. Top5 74. 2018-04-04 21:31:32,842 - Epoch: [89][ 450/ 500] Loss 2. Top1 51. Top5 74. 2018-04-04 21:31:35,338 - Epoch: [89][ 500/ 500] Loss 2. Top1 51. Top5 74. 2018-04-04 21:31:35,357 - ==> Top1: 51.838 Top5: 74.817 Loss: 2.150 2018-04-04 21:31:35,364 - Saving checkpoint 2018-04-04 21:31:39,251 - --- test --------------------- 2018-04-04 21:31:39,252 - 50000 samples (256 per mini-batch) 2018-04-04 21:31:51,512 - Test: [ 50/ 195] Loss 1. Top1 63. Top5 85. 2018-04-04 21:31:55,015 - Test: [ 100/ 195] Loss 1. Top1 60. Top5 83. 2018-04-04 21:31:58,732 - Test: [ 150/ 195] Loss 1. Top1 57. Top5 80. 2018-04-04 21:32:01,274 - ==> Top1: 56.606 Top5: 79.446 Loss: 1.893 

具体来看：

$ time python3 compress_classifier.py -a alexnet --lr 0.005 -p 50 ../../../data.imagenet -j 44 --epochs 90 --pretrained --compress=../sensitivity-pruning/alexnet.schedule_sensitivity.yaml 

在这个例子中，我们对预训练的AlexNet网络进行剪枝，并使用以下参数：

• 0.005的学习率
• 每50个mini-batch输出一次信息
• 使用44核线程加载数据（取决于个人服务器）
• 训练90个epoch
• 剪枝策略在alexnet.schedule_sensitivity.yaml中提供
• 输出到日志到logs

示例

Distiller附带了几个使用compress_classifier.py的示例，配置比较简单，可以将配置文件（YAML）直接作为命令行compress参数输入。具体可以浏览examples文件夹。同时网上附带了几个训练好的模型，不过大部分是剪枝方面的。

剪枝敏感度分析

训练后量化

Distiller支持训练模块的训练后量化，无需重新训练（使用基于范围的线性量化）。因此，任何模型（无论是否剪枝）都可以量化。要调用训练后量化，请使用--quantize-eval和--evaluate。 其他参数可用于控制量化参数：

Arguments controlling quantization at evaluation time("post-training quantization"): --quantize-eval, --qe Applicable only if --evaluate is also set --qe-mode QE_MODE, --qem QE_MODE Linear quantization mode. Choices: asym_s | asym_u | sym --qe-bits-acts NUM_BITS, --qeba NUM_BITS Number of bits for quantization of activations --qe-bits-wts NUM_BITS, --qebw NUM_BITS Number of bits for quantization of weights --qe-bits-accum NUM_BITS Number of bits for quantization of the accumulator --qe-clip-acts, --qeca Enable clipping of activations using min/max values averaging over batch --qe-no-clip-layers LAYER_NAME [LAYER_NAME ...], --qencl LAYER_NAME [LAYER_NAME ...] List of fully-qualified layer names for which not to clip activations. Applicable only if --qe-clip-acts is also set --qe-per-channel, --qepc Enable per-channel quantization of weights (per output channel) 

下面是一个resnet18的量化示例：

$ python3 compress_classifier.py -a resnet18 ../../../data.imagenet --pretrained --quantize-eval --evaluate 

进行量化的模型将被转储到运行目录中。 它将包含量化的模型参数（数据类型仍然是FP32，但值将是整数）。计算出的量化参数（比例和零点）也存储在每个量化层中。

总结

本文简单介绍了Distiller的安装和示例使用，如何将实验的需求结合到这里面还需要更深入地研究一下。