Day 21 : 模型优化 - 剪枝 Pruning

如果说可以让模型缩小10倍，精度还维持水准，这是什麽巫术?
延续 Day 20 的模型优化作法，本次再结合剪枝技术做到更轻量的模型效果。

什麽是剪枝 Pruning

剪枝 Pruning 将无关紧要的权重 (weight) 删除归零，在压缩时因为稀疏矩阵的特性，能明显缩小尺寸，可以压缩到原来 1/3。
如果经过剪枝再量化的模型，甚至可以缩小的原来 1/10 大小。

模型优化剪枝实作

Colab 支援。
本示范采用 Tensorflow 模型优化模组的 prune_low_magnitude() ，可以将 Keras 模型在训练期间将影响较小的权重修剪归零。
```
!pip install tensorflow\_model\_optimization
```

建立基本模型

我们的基本模型以训练後量化相同的基准模型进行优化，模型一样采用tf.keras.datasets.mnist，用CNN进行建模。
```
ACCURACY:
{'baseline Keras model': 0.9574999809265137}

MODEL_SIZE:
{'baseline h5': 98136} 
```

使用剪枝调整模型

进行剪枝，另外因为剪枝模型方法有增加一层包装层，摘要显示的参数会增加。

# Get the pruning method
prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude

# Compute end step to finish pruning after 2 epochs.
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

# Define pruning schedule.
pruning_params = {
    'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(
        initial_sparsity=0.50,
        final_sparsity=0.80,
        begin_step=0,
        end_step=end_step)
    }

# Pass in the trained baseline model
model_for_pruning = prune_low_magnitude(
    baseline_model, 
    **pruning_params
    )

# `prune_low_magnitude` requires a recompile.
model_for_pruning.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
    )

model_for_pruning.summary()

参数数量增加了。

观察剪枝前後的模型权重 weight 变化

剪枝前，有些微弱的权重。

重新训练模型。并在 Callback 增加 tfmot.sparsity.keras.UpdatePruningStep() 参数。

# Callback to update pruning wrappers at each step
callbacks=[tfmot.sparsity.keras.UpdatePruningStep()]

# Train and prune the model
model_for_pruning.fit(
    train_images, 
    train_labels,
    epochs=epochs, 
    validation_split=validation_split,
    callbacks=callbacks
    )

重新训练後已修剪，观察同一层的权重变化，许多不重要的权重已归零。

剪枝後移除包装层

剪枝之後，您可以用tfmot.sparsity.keras.strip_pruning()删除包装层以具有与基线模型相同的层和参数。
此方法也有助於保存模型并导出为*.tflite档案格式。
剪枝後尚未压缩的档案，模型档案大小与原先一致，这也挺合理的毕竟都还占着位子。
```
MODEL_SIZE:
{'baseline h5': 98136,
'pruned non quantized h5': 98136} 
```

模型压缩3倍术

剪枝後的模型再压缩。

压缩後档案大小约为原本1/3，这是因为剪枝後归零的权重可以更有效的压缩。

import tempfile
import zipfile

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


MODEL_SIZE['pruned non quantized h5'] = get_gzipped_model_size('pruned_model.h5')

MODEL_SIZE:
{'baseline h5': 98136,
'pruned non quantized h5': 25665}

模型压缩10倍术

现在尝试将已精剪枝後的模型再量化。

量化原本就会缩小约3倍，将剪枝模型压缩後再量化，与基本模型相比，这使模型大小减少了约为原本1/10，而且精度还能维持水准。

# 剪枝压缩後再量化模型
converter = tf.lite.TFLiteConverter.from_keras_model(baseline_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]

tflite_model = converter.convert()

with open('pruned_quantized.tflite', 'wb') as f:
    f.write(tflite_model)

MODEL_SIZE:
{'baseline h5': 98136,
 'pruned non quantized h5': 25665,
 'pruned quantized tflite': 8129}

 ACCURACY
 {'baseline Keras model': 0.9574999809265137,
 'pruned and quantized tflite': 0.9683,
 'pruned model h5': 0.9685999751091003}