tensorflow对已经训练好的模型进行剪枝（使用tensorflow-model-optimization）

1：训练一个模型

为了方便，我们使用MINIST来训练我们的模型，而MINIST这一数据集，即使简单的两层全连接也能得到一个很好的结果，所以我们构建一个简单的模型并训练它。

import tensorflow as tf
from tensorflow import keras
mnist = keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# 归一化预处理数据，即将数据转换为（0，1）的范围.
train_images = train_images / 255.0
test_images = test_images / 255.0

# 定义模型结构
model = keras.Sequential([
  keras.layers.InputLayer(input_shape=(28, 28)),
  keras.layers.Reshape(target_shape=(28, 28, 1)),
  keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation='relu'),
  keras.layers.MaxPooling2D(pool_size=(2, 2)),
  keras.layers.Flatten(),
  keras.layers.Dense(10)
])

# 编译模型，SparseCategoricalCrossentropy是交叉熵函数，是标签是非one-hot编码下使用的
#如果标签是one-hot编码，需要使用CategoricalCrossentropy
model.compile(optimizer='adam',
              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
#训练模型
model.fit(
  train_images,
  train_labels,
  epochs=4,
  validation_split=0.1,
)

Train on 54000 samples, validate on 6000 samples
Epoch 1/4
54000/54000 [==============================] - 10s 191us/sample - loss: 0.2867 - accuracy: 0.9196 - val_loss: 0.1144 - val_accuracy: 0.9717
Epoch 2/4
54000/54000 [==============================] - 10s 189us/sample - loss: 0.1123 - accuracy: 0.9676 - val_loss: 0.0808 - val_accuracy: 0.9790
Epoch 3/4
54000/54000 [==============================] - 10s 180us/sample - loss: 0.0814 - accuracy: 0.9765 - val_loss: 0.0674 - val_accuracy: 0.9825
Epoch 4/4
54000/54000 [==============================] - 10s 193us/sample - loss: 0.0670 - accuracy: 0.9806 - val_loss: 0.0727 - val_accuracy: 0.9817

model_result = model.evaluate(test_images,test_labels,verbose=0)
print(model_result)

[0.06889190595522524, 0.9784]

保存未剪枝的模型：

_, keras_file = tempfile.mkstemp('.h5',dir='./')
tf.keras.models.save_model(model, keras_file, include_optimizer=False)
print('Saved baseline model to:', keras_file)

如果想了解temfile.mkstemp，可以看我的tempfile.mkstemp 详解

2：对整个模型进行剪枝

2.1 安装tensorflow-model-optimization模块

pip install tensorflow-model-optimization

2.2 tensorflow-model-optimization中的函数介绍

2.2.1 PolynomialDecay函数

该函数位于：tensorflow-model-optimization.sparsity.keras.PolynomialDecay，返回一个
PruningSchedule对象。

PolynomialDecay(
    initial_sparsity,
    final_sparsity,
    begin_step,
    end_step,
    power=3,
    frequency=100,
)

1. 以稀疏性为initial_sparsity开始，到达到稀疏性为final_sparsity结束。多少稀疏就代表着多少权重将会消失（变成0）。
2. 以begin_step为开始，到end_step这个期间，每隔frequency的steps，就修剪一次模型。
3. power是多项式衰减系数，用于以下公式：

current_sparsity = final_sparsity + (initial_sparsity - final_sparsity)
      * (1 - (step - begin_step)/(end_step - begin_step)) ^ exponent

此公式用来计算此时的稀疏性，exponent即power，API文档对power的解释：Exponent to be used in the sparsity function，其中的step推测是运行过程中的那一刻的step。

其图像大概如下：

稀疏度逐渐在最后逐渐趋于平缓，直至达到目标稀疏度。

2.2.2 prune_low_magnitude函数

tensorflow-model-optimization.sparsity.keras.prune_low_magnitude函数
修改要在训练期间修剪的tf.keras图层或模型

prune_low_magnitude(
    to_prune,
    **parms,
)

其中函数中的参数列表有以下几个：

  params = {
      'pruning_schedule': pruning_schedule,
      'block_size': block_size,
      'block_pooling_type': block_pooling_type
  }

1. to_prune 是指要修建的模型或者层
2. pruning_schdeule是用来控制稀疏率变化的对象，这个是函数必须的参数
3. block_size, 可选参数，默认为（1，1），表示以2级权重张量表示的块稀疏模式的尺寸（高度，宽度）,没太明白这句话是什么意思,原文： The dimensions (height, weight) for the block sparse pattern in rank-2 weight tensors，但是下个参数的英文原文给了我启发，‘pool’让我想到了池化的过程，这个block_size应该是类似于池化窗口的东西，每次作用在（height,weight）大小的权重，并以“AVG”或者“MAX”的方式去得到这个窗口的结果。
4. block_pooling_type，可选参数，用于在块中合并权重的函数， 必须为“ AVG”或“ MAX”，默认为“AVG”。（The function to use to pool weights in the block. Must be ‘AVG’ or ‘MAX’）

当然如果你只想指定to_prune 和pruning_schdeule这两个参数，那么就有两种方式去传入参数，一是以字典的形式传入pruning_schdeule。

parms={'pruning_schdeule':pruning_schdeule对象}
prune_low_magnitude(
   to_prune,
   **parms
)

（这个pruning_schdeule对象可以通过PolynomialDecay生成。）
更简单一种就是你可以直接这样写：

prune_low_magnitude(
   to_prune,
   pruning_schdeule对象
)

3：模型剪枝的原理

剪枝原理较为简单，简单理解，即在训练过程中，在迭代一定次数后，便对net中的接近0的权重，都置为0，已达到对模型剪枝的作用，以此反复，直到net的参数达到目标的稀疏度。这样，模型训练完成后，模型里面大多数的weight皆为0，那么，当我们使用zip进行压缩时，模型便可以得到很大程度的压缩，且在推断过程中，减少了很多的计算量。
————————————————
版权声明：本文为CSDN博主「程序猿也可以很哲学」的原创文章，遵循CC 4.0
BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/qq_16564093/article/details/103511733

4：对整个模型进行剪枝

import tensorflow_model_optimization as tfmot
prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude
#PolynomialDecay方法定义一个具有多项式衰减功能的修剪计划，也就是说修剪过程中的稀疏度是变化的，网络参数逐渐减少，稀疏度逐渐提高。
pruning_schedule = tfmot.sparsity.keras.PolynomialDecay(initial_sparsity=0.50,
                                                               final_sparsity=0.80,
                                                               begin_step=0,
                                                               end_step=end_step)

import numpy as np
batch_size = 128
epochs = 2
validation_split = 0.1

num_images = train_images.shape[0] * (1 - validation_split)
end_step = np.ceil(num_images / batch_size).astype(np.int32) * epochs
#修改要在训练期间修剪的tf.keras层或模型，本例修剪的是整个模型的参数
model_for_pruning = prune_low_magnitude(model, pruning_schedule)
# 修剪后的模型需要重新编译
model_for_pruning.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model_for_pruning.summary()

Model: "sequential_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
prune_low_magnitude_reshape_ (None, 28, 28, 1)         1         
_________________________________________________________________
prune_low_magnitude_conv2d_1 (None, 26, 26, 12)        230       
_________________________________________________________________
prune_low_magnitude_max_pool (None, 13, 13, 12)        1         
_________________________________________________________________
prune_low_magnitude_flatten_ (None, 2028)              1         
_________________________________________________________________
prune_low_magnitude_dense_1  (None, 10)                40572     
=================================================================
Total params: 40,805
Trainable params: 20,410
Non-trainable params: 20,395

然后对比下，未进行剪枝的参数量：

model.summary()

Model: "sequential_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
reshape_1 (Reshape)          (None, 28, 28, 1)         0         
_________________________________________________________________
conv2d_1 (Conv2D)            (None, 26, 26, 12)        120       
_________________________________________________________________
max_pooling2d_1 (MaxPooling2 (None, 13, 13, 12)        0         
_________________________________________________________________
flatten_1 (Flatten)          (None, 2028)              0         
_________________________________________________________________
dense_1 (Dense)              (None, 10)                20290     
=================================================================
Total params: 20,410
Trainable params: 20,410
Non-trainable params: 0

通过对比，发现经过剪枝后的参数量多了Non-trainable params: 20,395，即多了20395个不可训练的参数。是tensorflow-model-optimization为网络中的每个权重添加的不可训练掩码，表示是否要修剪该权重，掩码为0或1。在修剪完模型后，我们需要使用strip_pruning来删除暂时添加的这些Non-trainable params。

5：训练评估新的模型

logdir = tempfile.mkdtemp()

#UpdatePruningStep()函数是必须的，如果不加就会报错。
#PruningSummaries()提供log以便跟踪进度以及调试
callbacks = [
  tfmot.sparsity.keras.UpdatePruningStep(),
  tfmot.sparsity.keras.PruningSummaries(log_dir=logdir),
]
 
model_for_pruning.fit(train_images, train_labels,
                  batch_size=batch_size, epochs=epochs, validation_split=validation_split,
                  callbacks=callbacks)

Train on 54000 samples, validate on 6000 samples
Epoch 1/2
54000/54000 [==============================] - 9s 164us/sample - loss: 0.0869 - accuracy: 0.9768 - val_loss: 0.1107 - val_accuracy: 0.9690
Epoch 2/2
54000/54000 [==============================] - 8s 143us/sample - loss: 0.0973 - accuracy: 0.9733 - val_loss: 0.0851 - val_accuracy: 0.9747

UpdatePruningStep回调，使其在训练过程中处理修剪更新。
PruningSummaries提供用于跟踪进度和调试的日志。

让我们看看剪枝前后的分类结果：

model_for_pruning_result = model_for_pruning.evaluate(
   test_images, test_labels, verbose=0)
print(model_result)
print(model_for_pruning_result)

[0.06889190595522524, 0.9784]
[0.0919371309094131, 0.9724]

准确率降低了约0.6%，损失值上升了大概上升了0.03。

6：创建比之前小三倍的模型

使用strip_pruning和应用标准压缩算法（例如通过gzip）。

6.1使用strip_pruning去除之前的不可训练权重,并保存模型。

model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)

_, pruned_keras_file = tempfile.mkstemp('.h5',dir='./')
tf.keras.models.save_model(model_for_export, pruned_keras_file, include_optimizer=False)
print('Saved pruned Keras model to:', pruned_keras_file)

6.1使用gzip对比剪枝前后文件大小

def get_gzipped_model_size(file):
  import os
  import zipfile

  _, zipped_file = tempfile.mkstemp('.zip')
  with zipfile.ZipFile(zipped_file, 'w', compression=zipfile.ZIP_DEFLATED) as f:
    f.write(file)

  return os.path.getsize(zipped_file)

print("Size of gzipped baseline Keras model: %.2f bytes" % (get_gzipped_model_size(keras_file)))
print("Size of gzipped pruned Keras model: %.2f bytes" % (get_gzipped_model_size(pruned_keras_file)))

Size of gzipped baseline Keras model: 78129.00 bytes
Size of gzipped pruned Keras model: 25708.00 bytes

发现了大概是变成了之前的1/3大小，而精度却没有太大的损失。