發(fā)布人：Yunlu Li 和 Artsiom Ablavatski

簡介

剪枝是 TensorFlow 模型優(yōu)化工具包 （TF MOT） 中提供的核心優(yōu)化技術(shù)之一。該技術(shù)不僅有助于減小模型尺寸，還可用于加速移動設(shè)備和網(wǎng)絡(luò)上的 CPU 處理。隨著現(xiàn)代計算密集型模型的出現(xiàn)，剪枝（作為一種模型優(yōu)化技術(shù)）領(lǐng)域已經(jīng)引起了極大關(guān)注。目前開發(fā)者可以輕松地對密集網(wǎng)絡(luò)進行剪枝（即將一部分權(quán)重設(shè)置為零），所導(dǎo)致的質(zhì)量下降可以忽略不計。

減小

https://arxiv.org/abs/1911.0972

今天，我們很高興地發(fā)布了對 TF MOT 剪枝 API 的一系列更新。這些更新將簡化剪枝技術(shù)，讓開發(fā)者能夠構(gòu)建稀疏模型，以便快速開展設(shè)備端推理。

剪枝 API

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

TF MOT 的更新

TensorFlow 通過 TensorFlow 模型優(yōu)化工具包 （TF MOT） 剪枝 API，為神經(jīng)網(wǎng)絡(luò)剪枝提供長期支持。該 API 于 2019 年推出，并引入了剪枝的基本基元，讓世界各地的研究人員都能使用這一新的優(yōu)化技術(shù)。如今，我們很高興地發(fā)布此 API 的實驗性更新，進一步推動剪枝模型發(fā)展。我們將會發(fā)布一些工具，來簡化剪枝的控制，并減少設(shè)備端推理延遲的時間。

TF MOT 剪枝 API 功能廣泛，為用戶提供了操縱模型的工具：

prune_low_magnitude 函數(shù)負(fù)責(zé)將 PruneLowMagnitude 封裝容器應(yīng)用至模型的每一層中

PruneLowMagnitude 封裝容器負(fù)責(zé)處理低層級的剪枝邏輯

PruningSchedule 負(fù)責(zé)控制應(yīng)用剪枝的時機

PruningSummaries 負(fù)責(zé)回調(diào)剪枝進程的日志

這些抽象概念便于用戶幾乎控制剪枝模型的任何方面（即，如何剪枝 （ PruneLowMagnitude ）、何時剪枝 （ PruningSchedule ），以及如何跟蹤剪枝的進程 （ PruningSummaries ））。但“ PruneLowMagnitude 封裝容器的應(yīng)用位置”這一剪枝內(nèi)容除外。如今發(fā)布了一個 TF MOT PruningPolicy 擴展程序，我們感到很高興。因為這個類可以用來控制 PruneLowMagnitude 封裝容器應(yīng)用于模型的部分。PruningPolicy 的實例則會被用作 prune_low_magnitude 函數(shù)的一個參數(shù)，包含以下功能：

通過 allow_pruning 函數(shù)來控制每個層上，用于應(yīng)用剪枝封裝容器的位置

通過 ensure_model_supports_pruning 函數(shù)檢查整個模型是否支持剪枝

PruningPolicy 是抽象的接口，根據(jù)特定應(yīng)用的不同而存在多種實施方式。為了通過 XNNPACK 優(yōu)化 CPU 的延遲時間，在具體實施時，PruneForLatencyOnXNNPack 將只對模型中可以通過設(shè)備端稀疏推理加速的部分應(yīng)用剪枝封裝容器，而對網(wǎng)絡(luò)的其他部分不作處理。這種選擇性的剪枝便于保證模型的質(zhì)量，同時關(guān)注模型中可以通過稀疏性加速的部分。

下面的例子展示了 PruneForLatencyOnXNNPack 策略在 MobileNetV2 上的應(yīng)用情況（完整示例請參閱最近推出的 colab）：

import tensorflow as tf

import tensorflow_model_optimization as tfmot

prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude

# 請參閱以下函數(shù)的實現(xiàn)實施情況。

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 模型上執(zhí)行 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’：

# 第一次卷積的默認(rèn)填充“SAME”與# XNNPACK 稀疏推理不兼容。

layer.padding = ‘valid’# 我們要求模型重新構(gòu)建，因為我們改變了填充參數(shù)。

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

執(zhí)行 Keras

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

PruneForLatencyOnXNNPack 策略只對具有 1x1 內(nèi)核大小的卷積應(yīng)用剪枝封裝容器，因為只有這些層可以通過使用 XNNPACK，讓 CPU 在速度上的提升高達 2 倍。該策略不會對其余層進行處理，這將便于網(wǎng)絡(luò)在剪枝步驟的質(zhì)量下降后得到恢復(fù)。

此外，該策略通過使用 ensure_model_supports_pruning 方法來驗證模型是否支持剪枝。在對稀疏模型進行訓(xùn)練和轉(zhuǎn)換后，我們建議在調(diào)試模式下使用 TensorFlow Lite 基準(zhǔn)實用程序，以便確認(rèn)最終的模型是否與 XNNPack 的稀疏推理后端兼容。

速度上的提升

https://arxiv.org/abs/1911.09723

實用程序

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

我們希望這個新引入的實驗性 API 可以發(fā)揮實際作用，并在未來繼續(xù)提高其穩(wěn)定性和靈活性。

改善壓縮和延遲時間

模型壓縮是對模型應(yīng)用剪枝的另一大好處。使用智能壓縮格式可以高效地存儲模型權(quán)重，從而顯著減小模型的大小。

格式

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

TFLite 采用 TACO 格式來編碼稀疏張量。與廣泛使用的格式（如 CSR 和 CSC）相比，TACO 格式具有以下幾個優(yōu)勢：

1. 支持靈活的遍歷順序，輕松存儲以行或列為主張量。

2. 支持如卷積算子的 4-D 過濾器等多維度稀疏張量。

3. 可以將塊結(jié)構(gòu)表示為張量的內(nèi)部維度（例如將 2x2 的內(nèi)部塊結(jié)構(gòu)表示為 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

我們還調(diào)整了格式，為存儲非零元素索引的元數(shù)據(jù)使用提供靈活的數(shù)據(jù)類型。對于小型張量或具有 int8_t 等緊湊數(shù)據(jù)類型的張量而言，這可以減少其存儲開銷。

靈活的數(shù)據(jù)類型

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

為了在模型轉(zhuǎn)換過程中實現(xiàn)實際大小的減小，需要應(yīng)用 tf.lite.Optimize.EXPERIMENTAL_SPARSITY 進行優(yōu)化。這項優(yōu)化可以對稀疏張量模型進行檢查，并將其轉(zhuǎn)換為高效的存儲格式。它還能與量化無縫銜接，您可以將兩者結(jié)合使用，進一步實現(xiàn)模型壓縮。這種轉(zhuǎn)換的完整示例如下：

# 從模型中移除剪枝封裝容器。

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

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

# 結(jié)合使用 float16 量化與稀疏優(yōu)化# 后者可緊湊地存儲經(jīng)過剪枝的權(quán)重。

converter.optimizations = ［

tf.lite.Optimize.EXPERIMENTAL_SPARSITY， # 啟用大小減小優(yōu)化。

tf.lite.Optimize.DEFAULT # 啟用轉(zhuǎn)換量化。

］

converter.target_spec.supported_types = ［tf.float16］

tflite_buffer = converter.convert（）

在應(yīng)用 tf.lite.Optimize.EXPERIMENTAL_SPARSITY 優(yōu)化和 PruneForLatencyOnXNNPack 剪枝策略后，可以將大小縮減 2 倍，如圖 1 所示：

除了減小大小，剪枝還可以通過 XNNPACK 為 CPU 提供推理加速。在 PruneForLatencyOnXNNPack 剪枝策略的助力下，我們在啟用 use_xnnpack 選項時，使用 TensorFlow Lite 基準(zhǔn)對 Pixel 4 上的 MobileNetV2 模型進行了 CPU 推理延遲時間的消減研究：

圖 2 中的這項研究表明，在使用 XNNPACK 的移動設(shè)備上運行時，延遲時間改善了 1.7x。訓(xùn)練稀疏 MobileNetV2 模型的策略以及超參數(shù)和預(yù)訓(xùn)練的檢查點已在 Elsen 等人的研究中有所描述。

Elsen 等人

https://arxiv.org/abs/1911.09723

剪枝技巧及提示

剪枝意識訓(xùn)練是模型優(yōu)化的關(guān)鍵步驟之一。訓(xùn)練中涉及許多超參數(shù)，其中一些參數(shù)（如剪枝計劃和學(xué)習(xí)率）會對模型的最終質(zhì)量產(chǎn)生巨大的影響。雖然已經(jīng)提出了許多策略，但一個簡單而有效的 3 步策略（參見表 1）在大多數(shù)用例中均展現(xiàn)出了強大的性能。該策略構(gòu)建在 Zhu & Gupta 的成熟方法之上，無需廣泛的再訓(xùn)練即可產(chǎn)生良好的效果：

Zhu & Gupta

https://arxiv.org/abs/1710.01878

這個策略所導(dǎo)致的訓(xùn)練時間大幅增加（約 3 倍）是不可避免的。然而，在搭配使用 PolynomialDecay 剪枝計劃的情況下，就經(jīng)過顯著剪枝 （》70%） 的神經(jīng)網(wǎng)絡(luò)而言，這個 3 步策略的效果下降是有限，有時甚至沒有下降。

MediaPipe 中的剪枝模型

隨著 TF MOT 剪枝 API 的更新，我們很激動地為一些 MediaPipe 解決方案發(fā)布剪枝模型。發(fā)布的模型包括姿勢和面部檢測器，以及剪枝后的手部跟蹤模型。所有這些模型都利用新引入的功能，通過 3 步剪枝策略進行了訓(xùn)練。與密集型基線相比，發(fā)布的剪枝模型在大小方面得到了明顯的減小，通過 XNNPACK 在 CPU 上運行時甚至呈現(xiàn)了出卓越的性能。在質(zhì)量方面，剪枝模型達到了類似的指標(biāo)，包括在我們的公平數(shù)據(jù)集上的評估（詳見模型卡）。解決方案的并行演示如下：

GPU 的剪枝

雖然在 GPU 上利用稀疏性可能具有挑戰(zhàn)性，但近期的研究已經(jīng)在提升這些平臺上的稀疏操作性能方面取得了進展。在流行的框架中添加對稀疏矩陣和稀疏操作的強力支持的勢頭已經(jīng)出現(xiàn)，一流的 GPU 最近也為某些形式的結(jié)構(gòu)化稀疏性增加了硬件加速支持。展望未來，在訓(xùn)練和推理中對支持稀疏性的軟硬件的改進將是該領(lǐng)域進步的一個關(guān)鍵因素。

對稀疏矩陣和稀疏操作的強力支持

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

未來方向

TF MOT 提供了多種模型優(yōu)化方法。經(jīng)過證明，其中許多方法都對高效的設(shè)備端模型推理至關(guān)重要。除了低量級剪枝算法外，我們將繼續(xù)擴展 TF MOT 剪枝 API，同時研究剪枝和量化技術(shù)的結(jié)合，來實現(xiàn)更好的設(shè)備端推理結(jié)果。敬請期待！

致謝

衷心感謝參與此項目的所有人員：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