知識蒸餾與量化技術最佳化模型效能

– 9 分鐘閱讀

模型最佳化是機器學習領域的重要議題，如何在兼顧效能的同時降低運算成本和資源消耗一直是開發者努力的方向。本文將探討知識蒸餾和量化技術的結合應用，展示如何有效提升模型效率。首先，我們會利用知識蒸餾技術訓練一個精簡的 DistilBERT 模型，並使用 Optuna 框架進行超引數搜尋，找出最佳的訓練引陣列合，例如訓練輪數、α 值和溫度等。接著，我們會介紹量化技術，將模型權重從 FP32 轉換為 INT8，以減少模型大小和運算成本。文章中將提供 PyTorch 程式碼範例，示範如何進行權重量化和動態量化，並比較量化前後的效能差異，包含模型大小、延遲和準確度等指標，驗證量化技術在實際應用中的有效性。

知識蒸餾最佳化模型效能

在前面的章節中，我們探討瞭如何利用知識蒸餾技術來訓練一個精簡的DistilBERT模型，以達到與教師模型（BERT）相當的效能。現在，我們將探討如何進一步最佳化這個過程，以獲得更好的結果。

使用Optuna進行超引數最佳化

為了找到最佳的超引陣列合，我們採用了Optuna這個強大的超引數最佳化框架。Optuna透過定義一個目標函式，並在多次試驗中對其進行最佳化，從而找到最佳的超引數。

import optuna

def objective(trial):
    x = trial.suggest_float("x", -2, 2)
    y = trial.suggest_float("y", -2, 2)
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2

study = optuna.create_study()
study.optimize(objective, n_trials=1000)

內容解密：

  1. Optuna的目標函式定義：我們定義了一個名為objective的函式，該函式接受一個trial物件作為輸入，並傳回一個需要被最佳化的目標值。在這個例子中，我們使用了Rosenbrock函式作為目標函式。
  2. 超引數建議：trial.suggest_float用於建議浮點數型別的超引數，範圍在-2到2之間。
  3. 建立研究物件：透過optuna.create_study()建立一個研究物件，用於管理多個試驗。
  4. 最佳化過程：study.optimize方法對目標函式進行最佳化，指定了試驗次數為1000。

將Optuna應用於Transformers

我們將Optuna應用於Transformers的知識蒸餾過程中，定義了需要最佳化的超引數空間，包括訓練輪數（num_train_epochs）、α（alpha）和溫度（temperature）。

def hp_space(trial):
    return {
        "num_train_epochs": trial.suggest_int("num_train_epochs", 5, 10),
        "alpha": trial.suggest_float("alpha", 0, 1),
        "temperature": trial.suggest_int("temperature", 2, 20)
    }

內容解密：

  1. 超引數空間定義：我們定義了一個名為hp_space的函式，該函式傳回一個字典，包含了需要被最佳化的超引數。
  2. 整數和浮點數超引數：分別使用suggest_int和suggest_float來建議整數和浮點數型別的超引數。

執行超引數搜尋

透過hyperparameter_search方法執行超引數搜尋，指定了試驗次數和最佳化方向。

best_run = distilbert_trainer.hyperparameter_search(
    n_trials=20, direction="maximize", hp_space=hp_space)

內容解密：

  1. 超引數搜尋：使用hyperparameter_search方法對定義的超引數空間進行搜尋。
  2. 最佳試驗結果：傳回一個BestRun物件，包含了最佳試驗的結果和對應的超引數。

訓練最終模型

根據最佳超引陣列合，更新訓練引數並執行最終的訓練。

for k, v in best_run.hyperparameters.items():
    setattr(student_training_args, k, v)

distil_trainer = DistillationTrainer(
    model_init=student_init, teacher_model=teacher_model, 
    args=student_training_args, train_dataset=clinc_enc['train'], 
    eval_dataset=clinc_enc['validation'], compute_metrics=compute_metrics, 
    tokenizer=student_tokenizer)

distil_trainer.train()

內容解密：

  1. 更新訓練引數：根據最佳超引陣列合更新訓練引數。
  2. 建立新的訓練器：使用更新後的訓練引數建立一個新的DistillationTrainer例項。
  3. 執行訓練：呼叫train方法執行最終的訓練。

評估最終模型的效能

透過建立一個pipeline並執行效能基準測試，來評估最終模型的效能。

distilled_ckpt = "transformersbook/distilbert-base-uncased-distilled-clinc"
pipe = pipeline("text-classification", model=distilled_ckpt)

optim_type = "Distillation"
pb = PerformanceBenchmark(pipe, clinc["test"], optim_type=optim_type)
perf_metrics.update(pb.run_benchmark())

內容解密：

  1. 建立pipeline：使用最終訓練好的模型建立一個文字分類別pipeline。
  2. 效能基準測試：執行效能基準測試，並更新效能指標。

結果分析

透過比較不同模型的效能指標，我們可以看到知識蒸餾技術顯著提高了模型的準確度，同時保持了較低的延遲。

透過量化技術加速模型運作

在前面的章節中，我們探討了知識蒸餾（Knowledge Distillation）如何將大型模型的知識轉移到較小的模型中，從而降低運算和記憶體成本。然而，我們還可以進一步最佳化模型的效能，透過一種稱為量化（Quantization）的技術來實作這一點。量化並不是減少計算次數，而是透過使用低精確度資料型別（如8位元整數，INT8）來取代常見的32位元浮點數（FP32），從而使運算更加高效。

浮點數與定點數簡介

大多數Transformer模型在預訓練和微調時使用浮點數（通常是FP32或FP16與FP32的混合），因為它們提供了足夠的精確度來處理權重、啟用值和梯度的不同範圍。浮點數表示法允許我們透過指數和尾數來表示非常廣泛的實數範圍。然而，在模型訓練完成後，我們只需要執行前向傳遞來進行推理，因此可以降低資料型別的精確度，而不會對準確率產生太大影響。

對於神經網路來說，常見的做法是使用定點格式來表示低精確度資料型別，即將實數表示為B位元整數，並透過一個共同的比例因子進行縮放。量化的基本思想是將浮點數值離散化，將它們對映到一個較小的定點數範圍內，並在之間線性分配所有值。

量化的數學表示

量化的過程可以透過以下公式來描述：

[ f = S(q - Z) ]

其中，(S) 是縮放因子，是一個正的浮點數；(Z) 是零點，具有與 (q) 相同的型別，對應於浮點數值 (f = 0) 的量化值。

程式碼範例：權重量化

import matplotlib.pyplot as plt
import torch

# 載入模型的狀態字典
state_dict = pipe.model.state_dict()
weights = state_dict["distilbert.transformer.layer.0.attention.out_lin.weight"]

# 繪製權重值的分佈圖
plt.hist(weights.flatten().numpy(), bins=250, range=(-0.3,0.3), edgecolor="C0")
plt.show()

# 計算量化所需的引數
zero_point = 0
scale = (weights.max() - weights.min()) / (127 - (-128))

# 量化權重
quantized_weights = (weights / scale + zero_point).clamp(-128, 127).round().char()
print(quantized_weights)

內容解密：

  1. 載入權重：首先，從預訓練的DistilBERT模型中載入權重值。
  2. 繪製分佈圖：使用matplotlib繪製權重值的直方圖，以觀察其分佈範圍。
  3. 計算量化引數：根據權重值的最大值和最小值，計算出量化的縮放因子（scale）和零點（zero_point）。
  4. 量化權重：將浮點數權重值對映到8位元整數空間，並進行四捨五入和截斷，以得到量化的權重。

使用PyTorch簡化量化過程

PyTorch提供了quantize_per_tensor()函式，可以簡化量化的過程。結合特定的量化資料型別（如torch.qint8），可以最佳化整數運算。

from torch import quantize_per_tensor

dtype = torch.qint8
quantized_weights = quantize_per_tensor(weights, scale, zero_point, dtype)
print(quantized_weights.int_repr())

內容解密：

  1. 指定量化資料型別：選擇torch.qint8作為量化後的資料型別。
  2. 執行量化：使用quantize_per_tensor()函式對權重進行量化。
  3. 輸出量化結果：列印預出量化後的整數表示。

量化對Transformer權重的影響

為了完成我們的分析，讓我們比較一下使用FP32和INT8值計算兩個權重張量乘法所需的時間。對於FP32張量，我們可以使用PyTorch的@運算子進行乘法：

%%timeit
weights @ weights
393 μs ± 3.84 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

對於量化張量，我們需要QFunctional包裝類別，以便可以使用特殊的torch.qint8資料型別執行操作：

from torch.nn.quantized import QFunctional
q_fn = QFunctional()

這個類別支援各種基本操作，例如加法，在我們的例子中，我們可以按照以下方式計時量化張量的乘法：

%%timeit
q_fn.mul(quantized_weights, quantized_weights)
23.3 μs ± 298 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

內容解密：

  1. 使用%%timeit魔法命令來測量程式碼執行的平均時間。
  2. weights @ weights是使用FP32精確度進行矩陣乘法的範例。
  3. q_fn.mul(quantized_weights, quantized_weights)是使用INT8精確度進行矩陣乘法的範例，展現了量化後的效能提升。
  4. QFunctional類別用於支援量化張量的運算。

與我們的FP32計算相比，使用INT8張量的速度幾乎快了100倍！透過使用專用的後端來高效執行量化運算子，可以獲得更大的效能提升。在本文撰寫時，PyTorch支援：

  • 具有AVX2支援或更高的x86 CPU
  • ARM CPU（通常出現在移動/嵌入式裝置中）

由於INT8數字的位數是FP32的四分之一，量化還可以將記憶體儲存需求減少最多四倍。在我們的簡單示例中，我們可以透過比較權重張量及其量化版本的底層儲存大小來驗證這一點，使用Tensor.storage()函式和Python的sys模組中的getsizeof()函式：

import sys
sys.getsizeof(weights.storage()) / sys.getsizeof(quantized_weights.storage())
3.999633833760527

內容解密：

  1. sys.getsizeof(weights.storage())用於取得原始權重張量的儲存大小。
  2. sys.getsizeof(quantized_weights.storage())用於取得量化後權重張量的儲存大小。
  3. 量化可以顯著減少模型的記憶體佔用。

對於全尺寸的Transformer，實際的壓縮率取決於哪些層被量化（如我們在下一節中將看到的，通常只有線性層被量化）。

那麼，量化的缺點是什麼？在模型的計算圖中的每個點引入小的擾動，可能會累積並影響模型的效能。有幾種量化模型的方法，它們都有優缺點。對於深度神經網路，通常有三種主要的量化方法：

動態量化

在訓練期間不進行任何更改，而是在推理期間執行適應性調整。與所有量化方法一樣，模型的權重在推理之前被轉換為INT8。除了權重之外，模型的啟用也被量化。這種方法是動態的，因為量化是在執行時進行的。這意味著所有的矩陣乘法都可以使用高度最佳化的INT8函式進行計算。

靜態量化

在推理之前，透過觀察代表性的資料樣本來預先計算量化的方案。這樣可以避免INT8和FP32值之間的轉換，從而加快計算速度。然而，它需要存取良好的資料樣本，並在Pipeline中引入額外的步驟。

量化感知訓練

在訓練期間透過“偽”量化FP32值來模擬量化的效果。這樣可以在訓練和推理期間提高模型的效能指標。

對於Transformer模型，動態量化是目前最好的方法。在較小的電腦視覺模型中，限制因素是啟用的記憶體頻寬，因此通常使用靜態量化（或在效能下降太大的情況下使用量化感知訓練）。

在PyTorch中實作動態量化非常簡單，可以透過一行程式碼完成：

from torch.quantization import quantize_dynamic
model_quantized = quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8)

內容解密：

  1. quantize_dynamic函式用於對模型進行動態量化。
  2. {nn.Linear}指定了需要量化的層型別。
  3. dtype=torch.qint8指定了目標精確度。

透過將模型透過基準測試並視覺化結果，我們可以看到量化對模型效能的影響：

pipe = pipeline("text-classification", model=model_quantized, tokenizer=tokenizer)
optim_type = "Distillation + quantization"
pb = PerformanceBenchmark(pipe, clinc["test"], optim_type=optim_type)
perf_metrics.update(pb.run_benchmark())

內容解密：

  1. 使用量化後的模型進行基準測試。
  2. PerformanceBenchmark類別用於評估模型的效能指標。

最終結果顯示，模型的尺寸減少到132.40 MB，平均延遲降低到12.54毫秒，測試集上的準確率為0.876。這些結果表明，量化可以在保持模型效能的同時顯著提高其效率。

玄貓 BlackCat

玄貓 BlackCat

技術愛好者，專注於分享程式開發、雲端技術與 AI 應用的心得體會。