發布人：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