1. 首頁 > 生活百科 > >

Pytorch模型量化

生活百科Pytorch模型量化

在深度學習中，量化指的是使用更少的bit來存儲原本以浮點數存儲的tensor，以及使用更少的bit來完成原本以浮點數完成的計算 。這么做的好處主要有如下幾點：

  • 更少的模型體積，接近4倍的減少；
  • 可以更快的計算 ， 由于更少的內存訪問和更快的int8計算，可以快2~4倍 。
一個量化后的模型，其部分或者全部的tensor操作會使用int類型來計算，而不是使用量化之前的float類型 。當然，量化還需要底層硬件支持，x86 CPU（支持AVX2）、ARM CPU、Google TPU、Nvidia Volta/Turing/Ampere、Qualcomm DSP這些主流硬件都對量化提供了支持 。
PyTorch對量化的支持目前有如下三種方式：
  • Post Training Dynamic Quantization：模型訓練完畢后的動態量化；
  • Post Training Static Quantization：模型訓練完畢后的靜態量化；
  • QAT (Quantization Aware Training)：模型訓練中開啟量化 。
在開始這三部分之前，先介紹下最基礎的Tensor的量化 。
Tensor的量化
量化：$$公式1：xq=round(\frac{x}{scale}+zero\_point)$$
反量化：$$公式2：x = (xq-zero\_point)*scale$$
式中，scale是縮放因子，zero_point是零基準，也就是fp32中的零在量化tensor中的值
為了實現量化，PyTorch 引入了能夠表示量化數據的Quantized Tensor，可以存儲 int8/uint8/int32類型的數據，并攜帶有scale、zero_point這些參數 。把一個標準的float Tensor轉換為量化Tensor的步驟如下：
import torchx = torch.randn(2, 2, dtype=torch.float32)# tensor([[ 0.9872, -1.6833],#         [-0.9345,  0.6531]])# 公式1(量化)：xq = round(x / scale + zero_point)# 使用給定的scale和 zero_point 來把一個float tensor轉化為 quantized tensorxq = torch.quantize_per_tensor(x, scale=0.5, zero_point=8, dtype=torch.quint8)# tensor([[ 1.0000, -1.5000],#         [-1.0000,  0.5000]], size=(2, 2), dtype=torch.quint8,#        quantization_scheme=torch.per_tensor_affine, scale=0.5, zero_point=8)print(xq.int_repr())  # 給定一個量化的張量 ， 返回一個以 uint8_t 作為數據類型的張量# tensor([[10,  5],#         [ 6,  9]], dtype=torch.uint8)# 公式2(反量化)：xdq = (xq - zero_point) * scale# 使用給定的scale和 zero_point 來把一個 quantized tensor 轉化為 float tensorxdq = xq.dequantize()# tensor([[ 1.0000, -1.5000],#         [-1.0000,  0.5000]])xdq和x的值已經出現了偏差的事實告訴了我們兩個道理：
  • 量化會有精度損失
  • 我們隨便選取的scale和zp太爛，選擇合適的scale和zp可以有效降低精度損失 。不信你把scale和zp分別換成scale = 0.0036, zero_point = 0試試
而在PyTorch中，選擇合適的scale和zp的工作就由各種observer來完成 。
Tensor的量化支持兩種模式：per tensor 和 per channel 。
  • Per tensor：是說一個tensor里的所有value按照同一種方式去scale和offset；
  • Per channel：是對于tensor的某一個維度（通常是channel的維度）上的值按照一種方式去scale和offset ， 也就是一個tensor里有多種不同的scale和offset的方式（組成一個vector），如此以來，在量化的時候相比per tensor的方式會引入更少的錯誤 。PyTorch目前支持conv2d()、conv3d()、linear()的per channel量化 。
在我們正式了解pytorch模型量化前我們再來檢查一下pytorch的官方量化是否能滿足我們的需求，如果不能，后面的都不需要看了
 靜態量化動態量化nn.linearYYnn.Conv1d/2d/3dYN (因為pytorch認為卷積參數來了個太小了，對卷積核進行量化會造成更多損失 ， 所以pytorch選擇不量化)nn.LSTMN(LSTM的好像又可以了，官方給出了一個例子，傳送門)Ynn.GRUNYnn.RNNCellNYnn.GRUCellNYnn.LSTMCellNYnn.EmbeddingBagY(激活在fp32)Ynn.EmbeddingYNnn.MultiheadAttentionNNActivations大部分支持不變 ， 計算停留在fp32中第二點：pytorch模型的動態量化只量化權重 ， 不量化偏置
Post Training Dynamic Quantization (訓練后動態量化)意思就是對訓練后的模型權重執行動態量化 ， 將浮點模型轉換為動態量化模型 ， 僅對模型權重進行量化 ， 偏置不會量化 。默認情況下，僅對 Linear 和 RNN 變體量化 (因為這些layer的參數量很大，收益更高) 。

推薦閱讀