通過 TensorFlow模型優(yōu)化工具包TF MOT剪枝API

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

簡(jiǎn)介

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

減小

https://arxiv.org/abs/1911.0972

今天，我們很高興地發(fā)布了對(duì) TF MOT 剪枝 API 的一系列更新。這些更新將簡(jiǎn)化剪枝技術(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 的實(shí)驗(yàn)性更新，進(jìn)一步推動(dòng)剪枝模型發(fā)展。我們將會(huì)發(fā)布一些工具，來簡(jiǎn)化剪枝的控制，并減少設(shè)備端推理延遲的時(shí)間。

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

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

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

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

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

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

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

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

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

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

import tensorflow as tf

import tensorflow_model_optimization as tfmot

prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude

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

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 需要對(duì)第一次卷積的填充操作進(jìn)行稍微修改：

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)建，因?yàn)槲覀兏淖兞颂畛鋮?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 策略只對(duì)具有 1x1 內(nèi)核大小的卷積應(yīng)用剪枝封裝容器，因?yàn)橹挥羞@些層可以通過使用 XNNPACK，讓 CPU 在速度上的提升高達(dá) 2 倍。該策略不會(huì)對(duì)其余層進(jìn)行處理，這將便于網(wǎng)絡(luò)在剪枝步驟的質(zhì)量下降后得到恢復(fù)。

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

速度上的提升

https://arxiv.org/abs/1911.09723

實(shí)用程序

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

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

改善壓縮和延遲時(shí)間

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

格式

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

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

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

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)整了格式，為存儲(chǔ)非零元素索引的元數(shù)據(jù)使用提供靈活的數(shù)據(jù)類型。對(duì)于小型張量或具有 int8_t 等緊湊數(shù)據(jù)類型的張量而言，這可以減少其存儲(chǔ)開銷。

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

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

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

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

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

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

# 結(jié)合使用 float16 量化與稀疏優(yōu)化# 后者可緊湊地存儲(chǔ)經(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 剪枝策略的助力下，我們?cè)趩⒂?use_xnnpack 選項(xiàng)時(shí)，使用 TensorFlow Lite 基準(zhǔn)對(duì) Pixel 4 上的 MobileNetV2 模型進(jìn)行了 CPU 推理延遲時(shí)間的消減研究：

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

Elsen 等人

https://arxiv.org/abs/1911.09723

剪枝技巧及提示

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

Zhu & Gupta

https://arxiv.org/abs/1710.01878

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

MediaPipe 中的剪枝模型

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

GPU 的剪枝

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

對(duì)稀疏矩陣和稀疏操作的強(qiáng)力支持

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

未來方向

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

致謝

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