TensorFlow官方发布剪枝优化工具：参数减少80%，精度几乎不变

图0：TensorFlow官方发布剪枝优化工具：参数减少80%，精度几乎不变

本文经AI新媒体量子位（公众号ID:QbitAI）授权转载，转载请联系出处。

去年TensorFlow官方推出了模型优化工具，最多能将模型尺寸减小4倍，运行速度提高3倍。

最近现又有一款新工具加入模型优化“豪华套餐”，这就是基于Keras的剪枝优化工具。

训练AI模型有时需要大量硬件资源，但不是每个人都有4个GPU的豪华配置，剪枝优化可以帮你缩小模型尺寸，以较小的代价进行推理。

什么是权重剪枝？

权重剪枝（Weight Pruning）优化，就是消除权重张量中不必要的值，减少神经网络层之间的连接数量，减少计算中涉及的参数，从而降低操作次数。

图1：TensorFlow官方发布剪枝优化工具：参数减少80%，精度几乎不变

这样做的好处是压缩了网络的存储空间，尤其是稀疏张量特别适合压缩。例如，经过处理可以将MNIST的90％稀疏度模型从12MB压缩到2MB。

此外，权重剪枝与量化（quantization）兼容，从而产生复合效益。通过训练后量化（post-training quantization），还能将剪枝后的模型从2MB进一步压缩到仅0.5MB 。

TensorFlow官方承诺，将来TensorFlow Lite会增加对稀疏表示和计算的支持，从而扩展运行内存的压缩优势，并释放性能提升。

优化效果

权重剪枝优化可以用于不同任务、不同类型的模型，从图像处理的CNN用于语音处理的RNN。下表显示了其中一些实验结果。

图2：TensorFlow官方发布剪枝优化工具：参数减少80%，精度几乎不变

以GNMT从德语翻译到英语的模型为例，原模型的BLEU为29.47。指定80%的稀疏度，经优化后，张量中的非零参数可以从211M压缩到44M，准确度基本没有损失。

使用方法

现在的权重剪枝API建立在Keras之上，因此开发者可以非常方便地将此技术应用于任何现有的Keras训练模型中。

开发者可以指定最终目标稀疏度（比如50%），以及执行剪枝的计划（比如2000步开始剪枝，在4000步时停止，并且每100步进行一次)，以及剪枝结构的可选配置。

import tensorflow_model_optimization as tfmot
model = build_your_model()
pruning_schedule = tfmot.sparsity.keras.PolynomialDecay(
initial_sparsity=0.0, final_sparsity=0.5,
begin_step=2000, end_step=4000)
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(model, pruning_schedule=pruning_schedule)
…
model_for_pruning.fit

图3：TensorFlow官方发布剪枝优化工具：参数减少80%，精度几乎不变