聊聊模型量化原理與代碼實現
什么是模型量化？
模型量化是一種模型壓縮技術，其核心思想是將原本使用浮點數表示的模型參數轉換為整數表示，以此來減少模型的存儲空間需求并加速計算過程。
具體來說，量化可以將原本使用float32類型的數據轉換為int8類型的數據，這意味著每個權重值占用的空間從32位減少到8位，不僅減少了模型的大小，也降低了計算所需的帶寬和計算資源。
為什么需要進行量化？
隨著深度學習技術在多個領域(如CV、NLP、語音等)的快速應用，模型的規模越來越大，復雜度也越來越高。這導致了模型在內存占用、計算資源以及能耗方面的需求也隨之增加。如需將這些復雜的模型部署在一些低成本的手機、平板等嵌入式設備中，往往難以滿足。
「模型量化應運而生，它可以在損失少量精度的前提下對模型進行壓縮，使得原本只能在高性能服務器或GPU上運行的大模型能夠在資源受限的嵌入式設備上運行。」
模型量化分類
根據映射函數是否是線性可以分為兩類——即線性量化和非線性量化，本文主要研究的是線性量化技術。
1.線性量化
線性量化的過程可以用以下數學表達式來表示：
其中，
- q表示原始的浮點數值（通常是Float32）。
- Z表示浮點數值的偏移量（通常稱為 Zero Point）。
- S表示浮點數值的縮放因子（通常稱為Scale）。
- Round(?)表示四舍五入近似取整的數學函數，也可以使用向上或向下取整。
根據參數Z是否為零可以將線性量化分為兩類——即對稱量化和非對稱量化。
(1) 對稱量化
對稱量化，即使用一個映射公式將輸入浮點數據映射到[-128,127]的范圍內，圖中-max(|Xf|)表示的是輸入數據的最小值，max(|Xf|)表示輸入數據的最大值。
對稱量化的一個核心即零點的處理，映射公式需要保證原始零點（即輸入浮點數中的0）在量化后依然對應于整數區間的0。總而言之，對稱量化通過映射關系將輸入數據映射在[-128,127]的范圍內，對于映射關系而言，我們需要求解的參數即Z和S。
在對稱量化中，r是用有符號的整型數值(int8)來表示的，此時Z=0，且q=0時恰好有r=0。S的計算公式如下：
其中，
- n表示用來表示該數值的位寬。
- max(|x|)表示數據集中所有樣本的絕對值的最大值。
(2) 非對稱量化
非對稱量化，即使用一個映射公式將輸入數據映射到[0,255]的范圍內，圖中min(Xf)表示的是輸入數據的最小值，max(Xf)表示輸入數據的最大值。
對稱量化通過映射關系將輸入數據映射在[0,255]的范圍內，對于映射關系而言，我們需要求解的參數即Z和S。
在非對稱量化中，r 是用有符號的整型數值(uint8)來表示的。可以取Z=min(x)，S的計算公式如下：
2.逐層量化、逐組量化和逐通道量化
根據量化的粒度（即共享量化參數的范圍），可以將量化方法分為逐層量化、逐組量化和逐通道量化。
- 逐層量化：以一個層為單位，整個層的所有權重上使用相同的縮放因子 S 和偏移量 Z 。
- 逐組量化：將權重按組劃分，每個group使用一組S和Z。
- 逐通道量化：以通道為單位，每個channel單獨使用一組S和Z。
當 group=1 時，逐組量化與逐層量化等價；當group=num_filters （即dw卷積）時，逐組量化逐通道量化等價。
3.在線量化與離線量化
根據激活值的量化方式，可以分為在線量化和離線量化兩種方法。這兩種方法的主要區別在于量化參數（縮放因子S和偏移量Z）是否在實際推理過程中動態計算。
- 在線量化：指在實際推理過程中，根據實際的激活值動態計算量化參數S和Z。
- 離線量化：離線量化是指提前確定好激活值的量化參數S和Z。這樣，在實際推理時就可以直接使用這些預計算好的參數，而不需要動態計算，從而提高了推理速度。
離線量化通常采用以下幾種方法來確定量化參數：
- 指數平滑法：將校準數據集送入模型，收集每個量化層的輸出特征圖，計算每個batch的S和Z值，并通過指數平滑法來更新S和Z值。
- 直方圖截斷法：在計算量化參數S和Z的過程中，考慮到有些特征圖可能會出現偏離較遠的奇異值，導致最大值非常大。可以采用直方圖截取的形式，比如拋棄最大的前 1% 數據，以前 1% 分界點的數值作為最大值來計算量化參數。
- KL 散度校準法:通過計算量化前后的兩個分布之間的 KL 散度（也稱為相對熵）來評估這兩個分布之間的差異，以搜索并選取KL散度最小的量化參數Z和S作為最終的結果。
4.比特量化
根據存儲一個權重元素所需的位數，可以將其分為8bit量化、4bit量化、2bit量化和1bit量化。
- 二進制神經網絡：即在運行時具有二進制權重和激活的神經網絡，以及在訓練時計算參數的梯度。
- 三元權重網絡：即權重約束為+1,0和-1的神經網絡。
- XNOR網絡：即過濾器和卷積層的輸入是二進制的。XNOR網絡主要使用二進制運算來近似卷積。
模型量化原理詳解
1.原理詳解
模型量化橋接定點和浮點，建立一種有效的數據映射關系。要弄懂模型量化的原理就要弄懂這種數據映射關系。浮點與定點數據的轉換公式如下：
其中，
- R表示輸入的浮點數據
- Q表示量化之后的定點數據
- Z表示零點（Zero Point）的數值
- S表示縮放因子（Scale）的數值
根據S和Z這兩個參數來確定這個映射關系。求解 S 和 Z 有很多種方法，這里列舉中其中的一種求解方式（MinMax）如下：
其中，
- max(R)表示輸入浮點數值的最大值。
- min(R)表示輸入浮點數值的最小值。
- max(Q)表示量化之后的整數數值的最大值（127/255）。
- min(Q)表示量化之后的整數數值的最小值（-128/0）。
每通道或每張量的權重用int8進行定點量化的可表示范圍為[-127，127]，且zero-point就是量化值0。
每張量的激活值或輸入值用int8進行定點量化的可表示范圍為[-128，127]，其zero-point在[-128，127]內依據公式求得。
2.具體案例
在這個案例中，我們將展示如何根據給定的激活值范圍 [-2.0, 6.0] 使用 int8 類型進行定點量化的過程。
步驟1: 計算量化尺度S和zero-point Z。
量化尺度S的計算公式為：
Zero-point Z的計算公式為：
代入給定的值：
- 激活值范圍 [-2.0, 6.0]，因此 max_val = 6.0 和 min_val = -2.0
- 定點量化值范圍 [-128, 127]，因此 quant_max = 127 和 quant_min = -128
計算得到：
步驟 2: 對激活值進行量化.
使用計算出的 S 和 Z 值對一個具體的激活值進行量化。假設有一個真實的激活值R = 0.28，則量化后的值 Q為：
代入S和Z的值：
模型量化實現步驟
模型量化具體的執行步驟如下所示：
- 在量化前，需要先統計出輸入數據（通常是權重或者激活值）中的最小值 min_value 和最大值 max_value。
- 根據模型的需求選擇合適的量化類型，常見的有對稱量化（int8）和非對稱量化（uint8）。
- 根據選擇的量化類型，計算量化參數 Z（Zero point）和 S（Scale）。
- 根據計算出的量化參數 Z 和 S，對模型執行量化操作，即將 FP32 數據轉換為 INT8 數據。
- 驗證量化后的模型性能是否滿足要求。如果不滿足，可以嘗試使用不同的方式計算 S 和 Z，然后重新執行量化操作。
Pytorch模型量化詳解
PyTorch提供了三種量化模型的方法，具體包括訓練后動態量化、訓練后靜態量化和訓練時量化。
1.訓練后動態量化
訓練后動態量化(Post Training Dynamic Quantization，PTDQ)是最簡單的量化形式，其中權重被提前量化，而激活在推理過程中被動態量化。這種方法用于模型執行時間由從內存加載權重而不是計算矩陣乘法所支配的情況，適合批量較小的LSTM和Transformer模型。步驟如下:
- 準備模型：將模型設置為評估模式 (model.eval())，對于需要動態量化的模型，通常不需要添加額外的量化或反量化模塊。
- 量化模型：使用 torch.quantization.quantize_dynamic() 函數來量化模型，這個函數會自動識別模型中適合動態量化的層，并將其轉換為量化版本。
2.訓練后靜態量化
訓練后靜態量化(Post-Training Static Quantization, PTQ)是最常用的量化形式，其中權重是提前量化的，并且基于在校準過程中觀察模型的行為來預先計算激活張量的比例因子和偏差。CNN是一個典型的用例，訓練后量化通常是在內存帶寬和計算節省都很重要的情況下進行的。訓練后量化的步驟如下：
- 準備模型：添加 QuantStub 和 DeQuantStub 模塊，以指定在何處顯式量化和反量化激活值。確保不重復使用模塊。將需要重新量化的任何操作轉換為模塊的模式。
- 融合操作：將諸如 conv + relu 或 conv + batchnorm + relu 之類的組合操作融合在一起，以提高模型的準確性和性能。
- 指定量化配置：例如選擇對稱或非對稱量化以及MinMax或L2Norm校準技術。
- 使用 torch.quantization.prepare() 函數來插入觀察模塊，以便在校準期間觀察激活張量。
- 使用校準數據集對模型執行校準操作。
- 使用 torch.quantization.convert() 函數來轉換模型。包括計算并存儲每個激活張量要使用的比例和偏差值，并替換關鍵算子的量化實現。
3.訓練時量化
在某些情況下，訓練后量化不能提供足夠的準確性，這時可以使用訓練時量化(Quantization-Aware Training，QAT)。步驟:
- 準備模型:添加 QuantStub 和 DeQuantStub 模塊，以指定在何處顯式量化和反量化激活值。確保不重復使用模塊。將需要重新量化的任何操作轉換為模塊的模式。
- 將諸如 conv + relu 或 conv + batchnorm + relu 之類的組合操作融合在一起，以提高模型的準確性和性能。
- 指定偽量化配置:例如選擇對稱或非對稱量化以及MinMax或L2Norm校準技術.
- 用 torch.quantization.prepare_qat() 函數來插入偽量化模塊，以便在訓練過程中模擬量化。
- 使用標準訓練流程訓練或微調模型。
- 使用 torch.quantization.convert() 函數來轉換模型，包括計算并存儲每個激活張量要使用的比例和偏差值，并替換關鍵算子的量化實現。
示例代碼
下面是一個簡單的示例，展示了如何使用PyTorch進行模型量化。
# 導入第三方的庫函數
import os
from io import open
import time
import torch
import torch.nn as nn
import torch.quantization
import torch.nn.functional as F
# 創建LSTM模型類
class LSTMModel(nn.Module):
"""整個網絡包含一個encoder, 一個recurrent模塊和一個decoder."""
def __init__(self, ntoken, ninp, nhid, nlayers, dropout=0.5):
super(LSTMModel, self).__init__()
# 預定義一些網絡層
self.drop = nn.Dropout(dropout)
# 嵌入層
self.encoder = nn.Embedding(ntoken, ninp)
# LSTM層
self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout)
# 線性層
self.decoder = nn.Linear(nhid, ntoken)
self.init_weights()
self.nhid = nhid
self.nlayers = nlayers
def init_weights(self):
'''
初始化模型權重
'''
initrange = 0.1
self.encoder.weight.data.uniform_(-initrange, initrange)
self.decoder.bias.data.zero_()
self.decoder.weight.data.uniform_(-initrange, initrange)
def forward(self, input, hidden):
'''
搭建網絡并執行前向推理
'''
emb = self.drop(self.encoder(input))
output, hidden = self.rnn(emb, hidden)
output = self.drop(output)
decoded = self.decoder(output)
return decoded, hidden
def init_hidden(self, bsz):
'''
初始化hidden層的權重
'''
weight = next(self.parameters())
return (weight.new_zeros(self.nlayers, bsz, self.nhid),
weight.new_zeros(self.nlayers, bsz, self.nhid))
# 創建一個詞典類，用來處理數據
# 構建詞匯表，包括詞到索引的映射和索引到詞的映射
class Dictionary(object):
def __init__(self):
self.word2idx = {}
self.idx2word = []
def add_word(self, word):
'''
在詞典中添加新的word
'''
if word not in self.word2idx:
self.idx2word.append(word)
self.word2idx[word] = len(self.idx2word) - 1
return self.word2idx[word]
def __len__(self):
'''
返回詞典的長度
'''
return len(self.idx2word)
# Corpus 類:處理文本數據，包括讀取文件、構建詞匯表和將文本轉換為索引序列
class Corpus(object):
def __init__(self, path):
self.dictionary = Dictionary()
# 分別獲取訓練集、驗證集和測試集
self.train = self.tokenize(os.path.join(path, 'train.txt'))
self.valid = self.tokenize(os.path.join(path, 'valid.txt'))
self.test = self.tokenize(os.path.join(path, 'test.txt'))
def tokenize(self, path):
"""對輸入的文件執行分詞操作"""
assert os.path.exists(path)
# 將新的單詞添加到詞典中
with open(path, 'r', encoding="utf8") as f:
for line in f:
words = line.split() + ['<eos>']
for word in words:
self.dictionary.add_word(word)
# 標記文件的內容
with open(path, 'r', encoding="utf8") as f:
idss = []
for line in f:
words = line.split() + ['<eos>']
ids = []
for word in words:
ids.append(self.dictionary.word2idx[word])
idss.append(torch.tensor(ids).type(torch.int64))
ids = torch.cat(idss)
return ids
# 設置模型的路徑
model_data_filepath = 'data/'
corpus = Corpus(model_data_filepath + 'wikitext-2')
ntokens = len(corpus.dictionary)
# 搭建網絡模型
model = LSTMModel(
ntoken = ntokens,
ninp = 512,
nhid = 256,
nlayers = 5,
)
# 加載預訓練的模型權重
model.load_state_dict(
torch.load(
model_data_filepath + 'word_language_model_quantize.pth',
map_location=torch.device('cpu')
)
)
# 將模型切換為推理模式，并打印整個模型
model.eval()
print(model)
# 獲取一個隨機的輸入數值
input_ = torch.randint(ntokens, (1, 1), dtype=torch.long)
hidden = model.init_hidden(1)
temperature = 1.0
num_words = 1000
# 遍歷數據集進行前向推理并將結果保存起來
with open(model_data_filepath + 'out.txt', 'w') as outf:
with torch.no_grad(): # no tracking history
for i in range(num_words):
output, hidden = model(input_, hidden)
word_weights = output.squeeze().div(temperature).exp().cpu()
word_idx = torch.multinomial(word_weights, 1)[0]
input_.fill_(word_idx)
word = corpus.dictionary.idx2word[word_idx]
outf.write(str(word.encode('utf-8')) + ('\n' if i % 20 == 19 else ' '))
if i % 100 == 0:
print('| Generated {}/{} words'.format(i, 1000))
with open(model_data_filepath + 'out.txt', 'r') as outf:
all_output = outf.read()
print(all_output)
bptt = 25
criterion = nn.CrossEntropyLoss()
eval_batch_size = 1
# 創建測試數據集
def batchify(data, bsz):
# 對測試數據集進行分塊
nbatch = data.size(0) // bsz
# 去掉多余的元素
data = data.narrow(0, 0, nbatch * bsz)
# 在bsz批處理中平均劃分數據
return data.view(bsz, -1).t().contiguous()
test_data = batchify(corpus.test, eval_batch_size)
# 獲取bath塊的輸入數據
def get_batch(source, i):
seq_len = min(bptt, len(source) - 1 - i)
data = source[i:i+seq_len]
target = source[i+1:i+1+seq_len].view(-1)
return data, target
def repackage_hidden(h):
"""
用新的張量把隱藏的狀態包裝起來，把它們從歷史中分離出來
"""
if isinstance(h, torch.Tensor):
return h.detach()
else:
return tuple(repackage_hidden(v) for v in h)
# 評估函數
def evaluate(model_, data_source):
# 打開評估模式
model_.eval()
total_loss = 0.
hidden = model_.init_hidden(eval_batch_size)
with torch.no_grad():
for i in range(0, data_source.size(0) - 1, bptt):
# 獲取測試數據
data, targets = get_batch(data_source, i)
# 執行前向推理
output, hidden = model_(data, hidden)
hidden = repackage_hidden(hidden)
output_flat = output.view(-1, ntokens)
# 獲取訓練loss
total_loss += len(data) * criterion(output_flat, targets).item()
return total_loss / (len(data_source) - 1)
# 初始化動態量化模塊
quantized_model = torch.quantization.quantize_dynamic(
model, {nn.LSTM, nn.Linear}, dtype=torch.qint8
)
print(quantized_model)
def print_size_of_model(model):
torch.save(model.state_dict(), "temp.p")
print('Size (MB):', os.path.getsize("temp.p")/1e6)
os.remove('temp.p')
print_size_of_model(model)
print_size_of_model(quantized_model)
torch.set_num_threads(1)
# 評估模型的運行時間
def time_model_evaluation(model, test_data):
s = time.time()
loss = evaluate(model, test_data)
elapsed = time.time() - s
print('''loss: {0:.3f}\nelapsed time (seconds): {1:.1f}'''.format(loss, elapsed))
time_model_evaluation(model, test_data)
time_model_evaluation(quantized_model, test_data)
