作者：Yunlu Li 和 Artsiom Ablavatski

简介

修剪是 TensorFlow 模型优化工具包 (TF MOT) 中提供的核心优化技术之一。它不仅有助于 大幅减小 模型大小，还可以用于 加速 移动设备和 Web 上的 CPU 推理。随着现代计算密集型模型的出现，修剪作为模型优化技术领域受到了广泛关注，证明了密集网络可以轻松修剪（即，将一小部分权重设置为零）而不会造成质量下降。今天，我们很高兴宣布对 TF MOT 修剪 API 进行一系列更新，这些更新简化了修剪并使开发人员能够构建用于快速设备上推理的稀疏模型。

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

# See the implementation of the function below.
model = load_model_for_pruning()

model_for_pruning = prune_low_magnitude(
    model, pruning_policy=tfmot.sparsity.keras.PruneForLatencyOnXNNPack())

为了遵守 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':
      # The default padding `SAME` for the first convolution is incompatible
      # with XNNPACK sparse inference.
      layer.padding = 'valid'
      # We ask the model to rebuild since we've changed the padding parameter.
      layer.built = False
    return layer

  return tf.keras.models.clone_model(model, clone_function=clone_fn)

PruneForLatencyOnXNNPack 策略仅将修剪包装器应用于内核大小为 1x1 的卷积，因为只有这些层才能通过 XNNPACK 在 CPU 上 加速 最多 2 倍。其余层保持不变，允许网络在修剪步骤中质量下降后恢复。此外，该策略通过使用 ensure_model_supports_pruning 方法验证模型是否适合修剪。一旦稀疏模型经过训练和转换，我们建议使用 TensorFlow Lite 基准测试 实用程序 的调试模式确认最终模型是否与 XNNPack 的稀疏推理后端兼容。

我们希望这个新推出的实验性 API 将在实践中有所帮助，我们将在未来继续提高其稳定性和灵活性。

压缩和延迟改进

模型压缩是将修剪应用于模型的另一个主要好处。使用智能压缩 格式 允许有效地存储模型权重，从而显着减小模型大小。

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

1. 它支持灵活的遍历顺序，以便轻松地以行主序或列主序格式存储张量。
2. 它支持多维稀疏张量，例如卷积运算的 4 维滤波器。
3. 它可以表示块结构作为张量的内部维度（示例，具有 2x2 内部块结构的 4x4 张量）。

我们还调整了该格式以使用 灵活的数据类型 来存储非零元素索引的元数据。这减少了小型张量或具有紧凑数据类型（如 int8_t）的张量的存储开销。

为了在模型转换过程中实际实现大小减小，需要应用 tf.lite.Optimize.EXPERIMENTAL_SPARSITY 优化。这种优化处理检查模型以寻找稀疏张量并将其转换为有效的存储格式。它还与量化无缝配合使用，您可以将它们组合起来实现更激进的模型压缩。下面显示了这种转换的完整示例

# Remove the pruning wrappers from the model. 
model = tfmot.sparsity.keras.strip_pruning(model)

converter = tf.lite.TFLiteConverter.from_keras_model(model)
# We apply float16 quantization together with sparsity optimization that
# compactly stores pruned weights.
converter.optimizations = [
    tf.lite.Optimize.EXPERIMENTAL_SPARSITY,  # Enables size reduction optimization. 
    tf.lite.Optimize.DEFAULT  # Enables quantization at conversion.
]
converter.target_spec.supported_types = [tf.float16]
tflite_buffer = converter.convert()

在应用 tf.lite.Optimize.EXPERIMENTAL_SPARSITY 优化以及 PruneForLatencyOnXNNPack 修剪策略后，可以实现约 2 倍 的大小减小，如 图 1 所示

图 1. MobileNetV2 模型大小（float32 和 float16 类型）的消融研究，使用 PruneForLatencyOnXNNPack 修剪策略在不同的稀疏性级别上进行。仅修剪了 1x1 卷积层，其余层保持密集。

除了大小减小之外，修剪还可以通过 XNNPACK 为 CPU 提供推理加速。使用 PruneForLatencyOnXNNPack 修剪策略，我们对 Pixel 4 上的 MobileNetV2 模型的 CPU 推理延迟进行了消融研究，使用 TensorFlow Lite 基准测试，并启用了 use_xnnpack 选项

图 2. Pixel 4 设备上 MobileNetV2 模型在不同稀疏性级别上的 CPU 推理速度的消融研究。

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

修剪技术和技巧

修剪感知训练是模型优化中的关键步骤。训练中涉及许多超参数，其中一些超参数（如修剪计划和学习率）会对模型的最终质量产生重大影响。尽管已经提出了许多策略，但简单而有效的 3 步策略（见表 1）在我们的大多数用例中都取得了良好的效果。该策略建立在 Zhu & Gupta 提出的经过良好验证的方法之上，并在没有大量重新训练的情况下产生良好的结果

该策略不可避免地会导致训练时间显着增加（约 3 倍）。但是，与 PolynomialDecay 修剪计划结合使用时，这种 3 步策略能够在显着修剪的（> 70%）神经网络中实现有限的质量下降或没有质量下降。

MediaPipe 中的修剪模型

我们很高兴地宣布，除了更新 TF MOT 剪枝 API 之外，我们还发布了一些经过剪枝的模型，这些模型适用于一些 MediaPipe 解决方案。发布的模型包括 姿势 和 人脸 检测器，以及 剪枝的手部跟踪 模型。所有这些模型都是使用新引入的功能和 3 步剪枝策略训练的。与密集基线相比，发布的剪枝模型在通过 XNNPACK 在 CPU 上运行时，已证明模型大小显着减少，同时性能也得到了提升。从质量上来说，剪枝模型在我们的公平性数据集上的评估中获得了类似的指标（有关详细信息，请参阅 模型卡片）。下面展示了这些解决方案的并排演示。

图 3. 人脸（顶部）和姿势（底部）检测端到端示例中密集模型（左侧）和稀疏模型（右侧）的比较

用于 GPU 的剪枝

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

未来方向

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。