Categories
程式開發

如何減小機器學習模型的大小


導讀:眾所周知,機器學習模型需要大量的計算、內存和功耗,這就給我們在實時推理或在計算資源有限的嵌入式設備上運行模型帶來了瓶頸。更複雜、更龐大的模型帶來的是更多的內存消耗,進而帶來更多的功耗。這就是模型優化要解決的問題。模型優化的技術有剪枝、量化等,我們今天翻譯並分享了 Amandeep Singh 的文章,講述瞭如何讓模型變得再小一些,以饗讀者。

機器學習模型變得越來越大,計算成本也越來越高。嵌入式設備的內存、計算能力和電池都受到限制。但我們可以對模型進行優化,使其在這些設備上能夠順利運行。通過減小模型的大小,我們減少了需要執行的操作數量,從而減少了計算量。較小的模型也很容易轉化為更少的內存使用,也就更節能。人們一定會認為,減少計算次數可以減少功耗,但相反,從內存訪問獲得的功耗比進行加法或乘法運算要高出 1000 倍左右。現在,既然沒有免費的午餐,也就是說,所有一切都是有代價的,因此,我們就會失去模型的正確率。記住,這些加速的措施並不是為了訓練模型,而是為了進行推理。

剪枝

剪枝(Pruning)就是刪除對輸出貢獻不大的多餘網絡連接。剪枝網絡的想法可以追溯到 20 世紀 90 年代,即“最優腦損傷”(Optimal Brain Damage)和“最優腦手術”(Optimal Brain Surgeon)。這些方法使用Hessians 來確定連接的重要性,這也使得它們不適用於深度網絡。剪枝方法使用迭代訓練技術,即訓練→ 剪枝→ 微調。剪枝後的微調恢復了網絡經剪枝後丟失的正確率。一種方法是使用L1/L2 範數對網絡中的權重進行排序,並去掉最後的x% 的權重。其他類型的方法也使用排序,使用神經元的平均激活,神經元在驗證集上的激活次數為零,還有許多其他創造性的方法。這種方法是由Han 等人在2015 年的論文中首創的。

如何減小機器學習模型的大小 1

神經網絡的剪枝。 Han 等人

更近一些的是 2019 年,Frankle 等人在論文《彩票假說》(The Lottery Ticket Hypothesis)中發現,在每個深度神經網絡中都存在一個子集,在同等數量的訓練下,該子集也具有同樣的正確率。這些結果適用於非結構化剪枝,即剪枝整個網絡,從而得到一個稀疏網絡。稀疏網絡在 GPU 上效率低下,因為它們的計算沒有結構。為了補修這一點,需要進行結構化剪枝,即對網絡的一部分進行剪枝,例如某一層或某一通道。Liu 等人發現,前面討論的彩票假說在這裡並不適用。相反,他們發現,在剪枝之後重新訓練網絡比微調更好。除了性能之外,稀疏網絡還有其他用途嗎?是的,正如 Ahmed 等人的論文所指出的那樣,稀疏網絡在噪聲輸入的情況下更具健壯性。在 TensorFlow(tensorflow_model_optimization 包)和 PyTorch(torch.nn.utils.prune)都支持剪枝。

要在 PyTorch 中使用剪枝,你可以從 torch.nn.utils.prune 中選擇一個技術類,或者實現 BasePruningMethod 的子類。

from torch.nn.utils import prune
tensor = torch.rand(2, 5)
pruner = prune.L1Unstructured(amount=0.7)
pruned_tensor = pruner.prune(tensor)

為了對模塊進行剪枝,我們可以使用 torch.nn.utils.prune 中給出的剪枝方法(基本上就是上述的類的包裝器),並指定哪個模塊要進行剪枝,甚至是該模塊的哪個參數。

conv_1 = nn.Conv(3, 1, 2)
prune.ln_structured(module=conv_1, name='weight', amount=5, n=2, dim=1)

這將使用剪枝後的結果替換參數權重,並添加一個參數 weight_orig 來存儲輸入的未剪枝版本。剪枝掩碼(pruning mask)存儲為 weight_mask,並作為模塊緩衝區保存。這些參數可以通過 module.named_parameters()module.named_buffers() 來檢查。為了實現迭代剪枝,我們可以只在下一次迭代中應用剪枝方法,這樣它就可以正常工作了,這是因為 PurningContainer 在處理最終掩碼的計算時,考慮到了之前使用 computer_mask 方法的剪枝。

量化

量化(Quantization)是為了限制一個權重可以取的可能值的數量,這將減少一個權重可以減少的內存,從而減小模型的大小。實現這一點的一種方法是,更改用於存儲權重的浮點數的位寬。以 32 位浮點數或 FP32 到 FP16、或 8 位定點數形式存儲的數字,越來越多地以 8 位整數的形式存儲。減少位寬具有以下許多優點:

  • 從 32 位轉換到 8 位,可以讓我們立即獲得 4 倍的內存優勢。
  • 較低的位寬還意味著,我們可以在寄存器 / 高速緩存中壓縮更多的數字,從而減少內存訪問,進而減少時間和功耗。
  • 整數計算總是比浮點計算要快。

之所以可行,是因為神經網絡對其權重的微小擾動是非常健壯的,我們可以很輕鬆地捨去它們,而不會對網絡的正確率產生太大的影響。此外,由於訓練中使用的正則化技術,權重並不包含在非常大的範圍內,因此我們不必使用過大的範圍,比如,對於32 位浮點數,取$-3.4 times 10^{38} $ 到$3.4 times 10^{38}$ 就可以了。例如,在下圖中,MoboileNet 中的權重值都非常接近於零。

如何減小機器學習模型的大小 2

MobileNetV1 的 10 層的權重分佈

一個量化方案是我們如何將實際權重轉換為量化權重,該方案的一個最基本的形式是線性縮放。假設我們要講範圍 $[r_{min},r_{max}]$ 的值轉換為 $[0,I_{max}]$ 的整數範圍,其中,$I_{max}$ 為 $2^B-1$,$B$ 是整數表示的位寬。因此,

$r=frac{r_{max}-r_{min}}{I_{max}-0}(q-z)=s(q-z)$

其中,$r$ 是權重的原始值,$s$ 是比例,$q$ 是量化值,$z$ 是映射到 0.0f 的值。這也稱為仿射變換(affine mapping)。由於 $q$ 為整數,因此對結果進行四捨五入。現在的問題是,我們如何選擇 $r_{min}$ 和 $r_{max}$。實現這一點的簡單方法是生成權重和激活的分佈,然後用量化分佈計算他們的 KL 散度(Kullback-Leibler divergence,縮寫為 KLD 或 KL divergences),並使用與原始值差異最小的那個。一種更為優雅的方法是使用偽量化(Fake Quantization),即,在訓練期間將量化感知層引入網絡。這個想法是由 Jacob 等人 提出的。

如何減小機器學習模型的大小 3

(a)普通卷積層;(b)增加偽量化單元的捲積層;(c)量化網絡的時延與正確率的比較。 Jacob 等人

在訓練時,偽量化節點計算權重和激活的範圍,並存儲它們的移動平均值。完成訓練後,我們用這個範圍來對網絡進行量化,以獲得更好的性能。

Rastegari 等人的關於異或(XOR)網絡的論文、Courbariaux 等人的關於三值(Ternary)網絡的論文、Zhu 等人的關於二值(Binary)網絡的論文中也探討了更大的位寬。在 PyTorch 1.3 中,引入了量化支持。為量化操作引入了三種新的數據類型:torch.quint8torch.qint8torch.qint32。它還提供了各種量化技術,包含在 torch.quantization 中。

  • 訓練後動態量化:將浮點權重替換為其動態量化版本。默認情況下,只對權重交大的層(即線性和 RNN 變體)進行權重量化。
quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.LSTM, nn.Linear}, dtype=torch.qint8
)
  • 訓練後靜態量化:靜態量化不僅可以將浮點數權重轉換為整數,還可以記錄激活的分佈情況,並用於確定推理時的量化比例。為了支持這種校準類型的量化,我們在模型的開頭和結尾分別添加了 QuantStubDeQuantStub。它涉及下面提到的步驟。
myModel = load_model(saved_model_dir + float_model_file).to('cpu')
# Fuse Conv, bn and relu
myModel.fuse_model()

# Specify quantization configuration
# Start with simple min/max range estimation and per-tensor 
# quantization of weights
myModel.qconfig = torch.quantization.default_qconfig

torch.quantization.prepare(myModel, inplace=True)

# Calibrate with the training set
evaluate(myModel, criterion, data_loader, 
            neval_batches=num_calibration_batches)

# Convert to quantized model
torch.quantization.convert(myModel, inplace=True)
  • 量化感知訓練:在訓練時使用偽量化模塊來存儲比例。為了啟用量化感知訓練,我們使用 qconfig 作為 get_default_qat_qconfig('fbgemm'),使用 prepare_qat 來代替 prepare。之後,就可以對模型進行訓練或微調,在訓練結束時,使用 與上述相同的 torch.quantization.convert 得到量化模型。

PyTorch 中的訓練後量化目前僅支持 CPU 上的操作。

有關詳細的代碼示例,請參閱 PyTorch 文檔。在 TensorFlow 方面,通過 將 optimizations 參數設置為 tf.lite.Optimize.OPTIMIZE_FOR_SIZE,可以使用 TFLite 的 tf.lite.TFLiteConverter API 進行量化。偽量化是通過 tf.contrib.quantize 包啟用的。

作者介紹:

Amandeep Singh,供職於 99acres.com 的高級軟件工程師。之前曾在三星諾伊達(印度)研究院 工作。廣泛研究設備上的機器學習和自然語言處理的問題。

原文鏈接:

https://amandeepsp.github.io/ml-model-compression-part1/