自从我们发布了 模型优化工具包 - 一套可以让开发者（无论新手还是老手）优化机器学习模型的技术 - 我们一直都在努力完成我们的 路线图，以添加更多新方法和工具。今天，我们很高兴分享新的权重剪枝 API。

权重剪枝

优化机器学习程序可以采取多种不同的形式。幸运的是，神经网络已被证明对旨在实现这一目标的不同转换具有很强的适应能力。

其中一类优化方法旨在通过删除神经网络层之间的连接（从而删除参数）来减少计算中涉及的参数和操作数量。

权重剪枝 API 建立在 Keras 之上，因此开发者可以轻松地将此技术应用于任何现有的 Keras 训练程序。该 API 将成为模型优化工具包的 GitHub 存储库 的一部分，与许多即将推出的优化技术一同提供。

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(...) 

什么是权重剪枝？

权重剪枝是指消除权重张量中不必要的数值。实际上，我们正在将神经网络参数的值设置为零，以移除我们估计为神经网络各层之间不必要连接的值。这在训练过程中完成，以允许神经网络适应这些变化。

为什么权重剪枝有用？

这项工作的直接好处是磁盘压缩：稀疏张量适合压缩。因此，通过对经过剪枝的 TensorFlow 检查点或已转换的 TensorFlow Lite 模型应用简单的文件压缩，我们可以减小模型的存储和/或传输尺寸。例如，在 教程 中，我们展示了如何将 MNIST 的 90% 稀疏模型从 12MB 压缩到 2MB。

此外，通过多次实验，我们发现权重剪枝与量化兼容，从而带来了复合效益。在同一教程中，我们展示了如何通过应用 训练后量化，将经过剪枝的模型从 2MB 进一步压缩到 0.5MB。

未来，TensorFlow Lite 将添加对稀疏表示和计算的一流支持，从而将压缩效益扩展到运行时内存，并释放性能提升，因为稀疏张量允许我们跳过涉及零值的不必要计算。

多个模型的成果

在我们的实验中，我们已经验证了这种技术可以成功地应用于不同类型的模型，跨越不同的任务，从基于卷积的神经网络图像处理到使用循环神经网络的语音处理。

不同模型和任务的稀疏性结果。

它是如何工作的？

我们的基于 Keras 的权重剪枝 API 使用了一种简单但广泛适用的 算法，该算法旨在根据训练过程中连接的大小来迭代地移除连接。从根本上说，需要指定一个最终目标稀疏度（例如 90%），以及执行剪枝的计划（例如，在第 2000 步开始剪枝，在第 10000 步停止，并且每 100 步执行一次），以及可选的剪枝结构配置（例如，应用于某些形状的单个值或值的块）。

无稀疏性张量的示例（左），1x1 块中的稀疏性（中），以及 1x2 块中的稀疏性（右）。

随着训练的进行，剪枝例程将被安排执行，消除（即设置为零）具有最小幅值的权重（即最接近零的权重），直到达到当前稀疏性目标。每次安排剪枝例程执行时，都会重新计算当前稀疏性目标，从 0% 开始，直到通过逐渐增加它（根据一个平滑的斜坡函数）在剪枝计划结束时达到最终目标稀疏度。

具有从第 0 步到第 100 步开始剪枝计划，以及 90% 的最终目标稀疏度的稀疏性斜坡函数的示例。

与计划一样，斜坡函数可以根据需要调整。例如，在某些情况下，在达到一定收敛水平后，可以安排训练过程从某个步骤开始，或者在训练程序的总训练步骤之前结束剪枝，以便在最终目标稀疏度水平上进一步微调系统。有关这些配置的更多详细信息，请参考我们的 教程 和 文档。

在训练过程结束时，与“已剪枝”Keras 层相对应的张量将包含根据该层的最终目标稀疏度设定的零。

应用于张量的剪枝动画。黑色单元格表示非零权重存在的位置。随着训练的进行，稀疏度会增加。

新的文档和 GitHub 存储库

如前所述，权重剪枝 API 将成为一个新的 GitHub 项目和存储库 的一部分，该项目和存储库的目标是让机器学习模型在执行和/或表示方面更有效。如果您对机器学习这个令人兴奋的领域感兴趣，或者只是想要优化模型的资源，这是一个值得关注的项目。

鉴于该领域的重要性，我们还在 tensorflow.org/model_optimization 下创建了一个新的子网站，其中包含相关文档和资源。我们鼓励您立即尝试一下，并欢迎您的反馈。关注此博客以获取未来的公告！

鸣谢：Raziel Alvarez、Pulkit Bhuwalka、Lawrence Chan、Alan Chiao、Tim Davis、Suyog Gupta、Jian Li、Yunlu Li、Sarah Sirajuddin、Daniel Situnayake、Suharsh Sivakumar。