通过 TensorFlow模型优化工具包TF MOT剪枝API

发布人：Yunlu Li 和 Artsiom Ablavatski

简介

剪枝是 TensorFlow 模型优化工具包 （TF MOT） 中提供的核心优化技术之一。该技术不仅有助于减小模型尺寸，还可用于加速移动设备和网络上的 CPU 处理。随着现代计算密集型模型的出现，剪枝（作为一种模型优化技术）领域已经引起了极大关注。目前开发者可以轻松地对密集网络进行剪枝（即将一部分权重设置为零），所导致的质量下降可以忽略不计。

减小

https://arxiv.org/abs/1911.0972

今天，我们很高兴地发布了对 TF MOT 剪枝 API 的一系列更新。这些更新将简化剪枝技术，让开发者能够构建稀疏模型，以便快速开展设备端推理。

剪枝 API

https://tensorflow.google.cn/model_optimization/guide/pruning

TF MOT 的更新

TensorFlow 通过 TensorFlow 模型优化工具包 （TF MOT） 剪枝 API，为神经网络剪枝提供长期支持。该 API 于 2019 年推出，并引入了剪枝的基本基元，让世界各地的研究人员都能使用这一新的优化技术。如今，我们很高兴地发布此 API 的实验性更新，进一步推动剪枝模型发展。我们将会发布一些工具，来简化剪枝的控制，并减少设备端推理延迟的时间。

TF MOT 剪枝 API 功能广泛，为用户提供了操纵模型的工具：

prune_low_magnitude 函数负责将 PruneLowMagnitude 封装容器应用至模型的每一层中

PruneLowMagnitude 封装容器负责处理低层级的剪枝逻辑

PruningSchedule 负责控制应用剪枝的时机

PruningSummaries 负责回调剪枝进程的日志

这些抽象概念便于用户几乎控制剪枝模型的任何方面（即，如何剪枝 （ PruneLowMagnitude ）、何时剪枝 （ PruningSchedule ），以及如何跟踪剪枝的进程 （ PruningSummaries ））。但“ PruneLowMagnitude 封装容器的应用位置”这一剪枝内容除外。如今发布了一个 TF MOT PruningPolicy 扩展程序，我们感到很高兴。因为这个类可以用来控制 PruneLowMagnitude 封装容器应用于模型的部分。PruningPolicy 的实例则会被用作 prune_low_magnitude 函数的一个参数，包含以下功能：

通过 allow_pruning 函数来控制每个层上，用于应用剪枝封装容器的位置

通过 ensure_model_supports_pruning 函数检查整个模型是否支持剪枝

PruningPolicy 是抽象的接口，根据特定应用的不同而存在多种实施方式。为了通过 XNNPACK 优化 CPU 的延迟时间，在具体实施时，PruneForLatencyOnXNNPack 将只对模型中可以通过设备端稀疏推理加速的部分应用剪枝封装容器，而对网络的其他部分不作处理。这种选择性的剪枝便于保证模型的质量，同时关注模型中可以通过稀疏性加速的部分。

下面的例子展示了 PruneForLatencyOnXNNPack 策略在 MobileNetV2 上的应用情况（完整示例请参阅最近推出的 colab）：

import tensorflow as tf

import tensorflow_model_optimization as tfmot

prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude

# 请参阅以下函数的实现实施情况。

model = load_model_for_pruning（）

model_for_pruning = prune_low_magnitude（

model， pruning_policy=tfmot.sparsity.keras.PruneForLatencyOnXNNPack（））

MobileNetV2

https://arxiv.org/abs/1801.04381

colab

https://tensorflow.google.cn/model_optimization/guide/pruning/pruning_for_on_device_inference

为了遵循 XNNPACK 稀疏推理的约束，在 MobileNetV2 模型上执行 Keras 需要对第一次卷积的填充操作进行稍微修改：

def load_model_for_pruning（）：

input_tensor = tf.keras.layers.Input（（224， 224， 3））

input_tensor = tf.keras.layers.ZeroPadding2D（1）（input_tensor）

model = tf.keras.applications.MobileNetV2（input_tensor=input_tensor）

def clone_fn（layer）：if layer.name == ‘Conv1’：

# 第一次卷积的默认填充“SAME”与# XNNPACK 稀疏推理不兼容。

layer.padding = ‘valid’# 我们要求模型重新构建，因为我们改变了填充参数。

layer.built = Falsereturn layer

return tf.keras.models.clone_model（model， clone_function=clone_fn）

约束

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/delegates/xnnpack/README.md#sparse-inference

执行 Keras

https://github.com/keras-team/keras-applications/blob/master/keras_applications/mobilenet_v2.py

PruneForLatencyOnXNNPack 策略只对具有 1x1 内核大小的卷积应用剪枝封装容器，因为只有这些层可以通过使用 XNNPACK，让 CPU 在速度上的提升高达 2 倍。该策略不会对其余层进行处理，这将便于网络在剪枝步骤的质量下降后得到恢复。

此外，该策略通过使用 ensure_model_supports_pruning 方法来验证模型是否支持剪枝。在对稀疏模型进行训练和转换后，我们建议在调试模式下使用 TensorFlow Lite 基准实用程序，以便确认最终的模型是否与 XNNPack 的稀疏推理后端兼容。

速度上的提升

https://arxiv.org/abs/1911.09723

实用程序

https://tensorflow.google.cn/lite/performance/measurement

我们希望这个新引入的实验性 API 可以发挥实际作用，并在未来继续提高其稳定性和灵活性。

改善压缩和延迟时间

模型压缩是对模型应用剪枝的另一大好处。使用智能压缩格式可以高效地存储模型权重，从而显著减小模型的大小。

格式

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/schema/schema.fbs#L94

TFLite 采用 TACO 格式来编码稀疏张量。与广泛使用的格式（如 CSR 和 CSC）相比，TACO 格式具有以下几个优势：

1. 支持灵活的遍历顺序，轻松存储以行或列为主张量。

2. 支持如卷积算子的 4-D 过滤器等多维度稀疏张量。

3. 可以将块结构表示为张量的内部维度（例如将 2x2 的内部块结构表示为 4x4 的张量）。

CSR

https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_matrix.html

CSC

https://scipy-lectures.org/advanced/scipy_sparse/csc_matrix.html

例如

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/testdata/sparse_tensor.json

我们还调整了格式，为存储非零元素索引的元数据使用提供灵活的数据类型。对于小型张量或具有 int8_t 等紧凑数据类型的张量而言，这可以减少其存储开销。

灵活的数据类型

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/schema/schema.fbs#L125-L135

为了在模型转换过程中实现实际大小的减小，需要应用 tf.lite.Optimize.EXPERIMENTAL_SPARSITY 进行优化。这项优化可以对稀疏张量模型进行检查，并将其转换为高效的存储格式。它还能与量化无缝衔接，您可以将两者结合使用，进一步实现模型压缩。这种转换的完整示例如下：

# 从模型中移除剪枝封装容器。

model = tfmot.sparsity.keras.strip_pruning（model）

converter = tf.lite.TFLiteConverter.from_keras_model（model）

# 结合使用 float16 量化与稀疏优化# 后者可紧凑地存储经过剪枝的权重。

converter.optimizations = ［

tf.lite.Optimize.EXPERIMENTAL_SPARSITY， # 启用大小减小优化。

tf.lite.Optimize.DEFAULT # 启用转换量化。

］

converter.target_spec.supported_types = ［tf.float16］

tflite_buffer = converter.convert（）

在应用 tf.lite.Optimize.EXPERIMENTAL_SPARSITY 优化和 PruneForLatencyOnXNNPack 剪枝策略后，可以将大小缩减 2 倍，如图 1 所示：

除了减小大小，剪枝还可以通过 XNNPACK 为 CPU 提供推理加速。在 PruneForLatencyOnXNNPack 剪枝策略的助力下，我们在启用 use_xnnpack 选项时，使用 TensorFlow Lite 基准对 Pixel 4 上的 MobileNetV2 模型进行了 CPU 推理延迟时间的消减研究：

图 2 中的这项研究表明，在使用 XNNPACK 的移动设备上运行时，延迟时间改善了 1.7x。训练稀疏 MobileNetV2 模型的策略以及超参数和预训练的检查点已在 Elsen 等人的研究中有所描述。

Elsen 等人

https://arxiv.org/abs/1911.09723

剪枝技巧及提示

剪枝意识训练是模型优化的关键步骤之一。训练中涉及许多超参数，其中一些参数（如剪枝计划和学习率）会对模型的最终质量产生巨大的影响。虽然已经提出了许多策略，但一个简单而有效的 3 步策略（参见表 1）在大多数用例中均展现出了强大的性能。该策略构建在 Zhu & Gupta 的成熟方法之上，无需广泛的再训练即可产生良好的效果：

Zhu & Gupta

https://arxiv.org/abs/1710.01878

这个策略所导致的训练时间大幅增加（约 3 倍）是不可避免的。然而，在搭配使用 PolynomialDecay 剪枝计划的情况下，就经过显著剪枝 （》70%） 的神经网络而言，这个 3 步策略的效果下降是有限，有时甚至没有下降。

MediaPipe 中的剪枝模型

随着 TF MOT 剪枝 API 的更新，我们很激动地为一些 MediaPipe 解决方案发布剪枝模型。发布的模型包括姿势和面部检测器，以及剪枝后的手部跟踪模型。所有这些模型都利用新引入的功能，通过 3 步剪枝策略进行了训练。与密集型基线相比，发布的剪枝模型在大小方面得到了明显的减小，通过 XNNPACK 在 CPU 上运行时甚至呈现了出卓越的性能。在质量方面，剪枝模型达到了类似的指标，包括在我们的公平数据集上的评估（详见模型卡）。解决方案的并行演示如下：

GPU 的剪枝

虽然在 GPU 上利用稀疏性可能具有挑战性，但近期的研究已经在提升这些平台上的稀疏操作性能方面取得了进展。在流行的框架中添加对稀疏矩阵和稀疏操作的强力支持的势头已经出现，一流的 GPU 最近也为某些形式的结构化稀疏性增加了硬件加速支持。展望未来，在训练和推理中对支持稀疏性的软硬件的改进将是该领域进步的一个关键因素。

对稀疏矩阵和稀疏操作的强力支持

https://github.com/tensorflow/community/blob/master/rfcs/20200519-csr-sparse-matrix.md

未来方向

TF MOT 提供了多种模型优化方法。经过证明，其中许多方法都对高效的设备端模型推理至关重要。除了低量级剪枝算法外，我们将继续扩展 TF MOT 剪枝 API，同时研究剪枝和量化技术的结合，来实现更好的设备端推理结果。敬请期待！

致谢

衷心感谢参与此项目的所有人员：Karthik Raveendran、Ethan Kim、Marat Dukhan、Trevor Gale、Utku Evci、Erich Elsen、Frank Barchard、Yury Kartynnik、Valentin Bazarevsky、Matsvei Zhdanovich、Juhyun Lee、Chuo-Ling Chang、Ming Guang Yong、Jared Duke 和 Matthias Grundmann。

编辑：jq