[[265358]]

本文經(jīng)AI新媒體量子位（公眾號(hào)ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

去年TensorFlow官方推出了模型優(yōu)化工具，最多能將模型尺寸減小4倍，運(yùn)行速度提高3倍。

最近現(xiàn)又有一款新工具加入模型優(yōu)化“豪華套餐”，這就是基于Keras的剪枝優(yōu)化工具。

訓(xùn)練AI模型有時(shí)需要大量硬件資源，但不是每個(gè)人都有4個(gè)GPU的豪華配置，剪枝優(yōu)化可以幫你縮小模型尺寸，以較小的代價(jià)進(jìn)行推理。

什么是權(quán)重剪枝？

權(quán)重剪枝（Weight Pruning）優(yōu)化，就是消除權(quán)重張量中不必要的值，減少神經(jīng)網(wǎng)絡(luò)層之間的連接數(shù)量，減少計(jì)算中涉及的參數(shù)，從而降低操作次數(shù)。

這樣做的好處是壓縮了網(wǎng)絡(luò)的存儲(chǔ)空間，尤其是稀疏張量特別適合壓縮。例如，經(jīng)過(guò)處理可以將MNIST的90％稀疏度模型從12MB壓縮到2MB。

此外，權(quán)重剪枝與量化（quantization）兼容，從而產(chǎn)生復(fù)合效益。通過(guò)訓(xùn)練后量化（post-training quantization），還能將剪枝后的模型從2MB進(jìn)一步壓縮到僅0.5MB 。

TensorFlow官方承諾，將來(lái)TensorFlow Lite會(huì)增加對(duì)稀疏表示和計(jì)算的支持，從而擴(kuò)展運(yùn)行內(nèi)存的壓縮優(yōu)勢(shì)，并釋放性能提升。

優(yōu)化效果

權(quán)重剪枝優(yōu)化可以用于不同任務(wù)、不同類(lèi)型的模型，從圖像處理的CNN用于語(yǔ)音處理的RNN。下表顯示了其中一些實(shí)驗(yàn)結(jié)果。

以GNMT從德語(yǔ)翻譯到英語(yǔ)的模型為例，原模型的BLEU為29.47。指定80%的稀疏度，經(jīng)優(yōu)化后，張量中的非零參數(shù)可以從211M壓縮到44M，準(zhǔn)確度基本沒(méi)有損失。

使用方法

現(xiàn)在的權(quán)重剪枝API建立在Keras之上，因此開(kāi)發(fā)者可以非常方便地將此技術(shù)應(yīng)用于任何現(xiàn)有的Keras訓(xùn)練模型中。

開(kāi)發(fā)者可以指定最終目標(biāo)稀疏度（比如50%），以及執(zhí)行剪枝的計(jì)劃（比如2000步開(kāi)始剪枝，在4000步時(shí)停止，并且每100步進(jìn)行一次)，以及剪枝結(jié)構(gòu)的可選配置。

import tensorflow_model_optimization as tfmot 
model = build_your_model()  
pruning_schedule = tfmot.sparsity.keras.PolynomialDecay( 
 initial_sparsity=0.0, final_sparsity=0.5, 
 begin_step=2000, end_step=4000) 
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(model, pruning_schedule=pruning_schedule) 
...  
model_for_pruning.fit 

隨著訓(xùn)練的進(jìn)行，剪枝過(guò)程開(kāi)始被執(zhí)行。在這個(gè)過(guò)程中，它會(huì)消除消除張量中最接近零的權(quán)重，直到達(dá)到當(dāng)前稀疏度目標(biāo)。

每次計(jì)劃執(zhí)行剪枝程序時(shí)，都會(huì)重新計(jì)算當(dāng)前稀疏度目標(biāo)，根據(jù)平滑上升函數(shù)逐漸增加稀疏度來(lái)達(dá)到最終目標(biāo)稀疏度，從0%開(kāi)始直到結(jié)束。

用戶(hù)也可以根據(jù)需要調(diào)整這個(gè)上升函數(shù)。在某些情況下，可以安排訓(xùn)練過(guò)程在某個(gè)步驟達(dá)到一定收斂級(jí)別之后才開(kāi)始優(yōu)化，或者在訓(xùn)練總步數(shù)之前結(jié)束剪枝，以便在達(dá)到最終目標(biāo)稀疏度時(shí)進(jìn)一步微調(diào)系統(tǒng)。

GitHub地址：

https://github.com/tensorflow/model-optimization

官方教程：

https://www.tensorflow.org/model_optimization/guide/pruning/pruning_with_keras