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,
)
- 以稀疏性为initial_sparsity开始,到达到稀疏性为final_sparsity结束。多少稀疏就代表着多少权重将会消失(变成0)。
- 以begin_step为开始,到end_step这个期间,每隔frequency的steps,就修剪一次模型。
- 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
}
- to_prune 是指要修建的模型或者层
- pruning_schdeule是用来控制稀疏率变化的对象,这个是函数必须的参数
- 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”的方式去得到这个窗口的结果。
- 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大小,而精度却没有太大的损失。