Transformer 模型非常強大,但往往太大太慢,不適合實時應用。為了解決這個問題,我們來看看三種關鍵的優化技術:知識蒸餾、量化和ONNX 圖優化。這些技術可以顯著減少推理時間和內存使用。
為了說明每種技術的利弊,我們以意圖檢測為例,因為它是基于文本的助手的重要組成部分,實時對話中低延遲至關重要。
為生產環境基準測試 Transformer 模型
就像任何機器學習模型一樣,將 Transformer 部署到生產環境中,不僅僅是關于準確性——而是要在相互競爭的系統級需求之間做出明智的權衡。三個關鍵約束條件始終浮出水面:
- 模型性能:模型對真實世界數據的泛化能力如何?在錯誤代價高昂的領域——無論是由于監管風險、用戶影響還是規模——即使是精度或召回率的微小改進,也可能產生巨大的下游效應。在高風險場景中,引入“人在回路”可以幫助最小化關鍵錯誤。
- 延遲:模型返回預測的速度有多快?在實時應用中,尤其是在大規模運行的應用中,低延遲推理至關重要。例如,Stack Overflow 需要一個響應式的分類器,以即時標記有問題的評論,而不干擾用戶流程。
- 內存效率:我們如何部署像 GPT-2 或 T5 這樣擁有數十億參數且需要大量計算資源的模型?在移動設備或邊緣環境中部署時,內存成為一個關鍵約束,因為這些環境可能無法訪問高性能的云基礎設施,或者根本不存在。
為什么基準測試很重要:關鍵要點
在 Transformer 部署中未能平衡性能、延遲和內存約束可能導致:
- 用戶體驗下降:緩慢、無響應的模型會讓用戶感到沮喪,降低產品價值。
- 不必要的基礎設施成本:在流量很少的情況下在云服務器上運行大型模型,會導致過高的計算賬單和資源利用不足。
解決方案:構建目標基準
為了解決這些挑戰,我們將設計一個輕量級基準測試框架,它將:
- 評估核心約束(性能、延遲、內存)
- 在定義好的管道和測試集上運行
- 為應用模型優化技術(如量化、剪枝和蒸餾)奠定基礎
這從一個簡單且可擴展的基準測試類開始——這是系統性能分析和壓縮實驗的基礎。
class PerformanceBenchmark:def __init__(self, pipeline, dataset, optim_type="BERT 基線"):self.pipeline = pipelineself.dataset = datasetself.optim_type = optim_typedef compute_accuracy(self):passdef compute_size(self):passdef time_pipeline(self):pass
def run_benchmark(self):metrics = {}metrics[self.optim_type] = self.compute_size()metrics[self.optim_type].update(self.time_pipeline())metrics[self.optim_type].update(self.compute_accuracy())return metrics
用真實數據測量模型準確性
有了我們的基準測試框架后,是時候讓它“活”起來,通過計算模型在代表性測試集上的準確性。
為此,我們將使用CLINC150 數據集——這是一個廣泛用于意圖分類任務的基準數據集。這個數據集也被用來微調我們的基線 Transformer 模型,使其成為評估的理想起點。
from datasets import load_dataset
clinc = load_dataset("clinc_oos", "plus")
clinc
DatasetDict({train: Dataset({features: ['text', 'intent'],num_rows: 15250})validation: Dataset({features: ['text', 'intent'],num_rows: 3100})test: Dataset({features: ['text', 'intent'],num_rows: 5500})
})
了解 CLINC150 數據集結構
CLINC150 數據集中的每個條目包含:
- 一個用戶查詢(存儲在
text
字段中) - 其對應的意圖標簽(存儲在
intent
字段中)
為了基準測試,我們將關注測試集,因為它最能模擬真實世界的使用。為了了解數據格式,我們來檢查測試集中的一個樣本:
了解 CLINC150 數據集結構
CLINC150 數據集中的每個條目包含:一個用戶查詢(存儲在 text 字段中)其對應的意圖標簽(存儲在 intent 字段中)為了基準測試,我們將關注測試集,因為它最能模擬真實世界的使用。為了了解數據格式,我們來檢查測試集中的一個樣本:
意圖是以 ID 形式提供的,但我們可以輕松地通過訪問 Dataset.features 屬性來獲取字符串與 ID 之間的映射(反之亦然):
intents = clinc["test"].features["intent"]
intents.int2str(clinc["test"][42]["intent"])
'transfer'
現在我們已經對 CLINC150 數據集的內容有了基本的了解,讓我們來實現 compute_accuracy
函數。
from datasets import load_metric
accuracy_score = load_metric('accuracy')
accuracy_score
Metric(name: "accuracy", features: {'predictions': Value(dtype='int32',> id=None), 'references': Value(dtype='int32', id=None)}, usage: """
Args:predictions: Predicted labels, as returned by a model.references: Ground truth labels.normalize: If False, return the number of correctly classified samples.Otherwise, return the fraction of correctly classified samples.sample_weight: Sample weights.
Returns:accuracy: Accuracy score.
""", stored examples: 0)
為了評估我們模型的性能,我們將使用準確率指標——但它需要預測標簽和真實標簽都以整數 ID的形式表示。
以下是步驟:
- 生成預測:使用預訓練的管道對
text
字段進行預測。 - 將預測標簽轉換為整數 ID:使用
ClassLabel.str2int()
,它將字符串類名映射到它們對應的數值索引。 - 將所有預測和真實標簽分別收集到單獨的列表中。
- 計算準確率:將兩個列表傳遞給指標函數。
讓我們將這個邏輯集成到我們的 PerformanceBenchmark
類中,以自動化這個過程:
ef compute_accuracy(self):preds, labels = [], []for example in self.dataset:pred = self.pipeline(example["text"])[0]["label"]label = example["intent"]preds.append(intents.str2int(pred))labels.append(label)accuracy = accuracy_score.compute(predictions=preds, references=labels)print(f"測試集上的準確率 - {accuracy['accuracy']:.3f}")return accuracy
PerformanceBenchmark.compute_accuracy = compute_accuracy
為了了解我們模型的內存占用,我們將把它序列化到磁盤并測量其大小。PyTorch 提供了一個方便的方法來實現這一點,使用 torch.save
,它依賴于 Python 內置的 pickle
模塊。它可以用來持久化從模型和張量到普通 Python 對象的一切。
在 PyTorch 中保存模型時,推薦的方法是保存其 state_dict
——這是一個包含模型每一層的所有可學習參數(如權重和偏置)的字典。
讓我們看看我們基線 Transformer 模型的 state_dict
里有什么:
list(pipe.model.state_dict().items())[42]
('bert.encoder.layer.2.attention.self.value.weight',tensor([[-1.0526e-02, -3.2215e-02, 2.2097e-02, ..., -6.0953e-03,4.6521e-03, 2.9844e-02],[-1.4964e-02, -1.0915e-02, 5.2396e-04, ..., 3.2047e-05,-2.6890e-02, -2.1943e-02],[-2.9640e-02, -3.7842e-03, -1.2582e-02, ..., -1.0917e-02,3.1152e-02, -9.7786e-03],...,[-1.5116e-02, -3.3226e-02, 4.2063e-02, ..., -5.2652e-03,1.1093e-02, 2.9703e-03],[-3.6809e-02, 5.6848e-02, -2.6544e-02, ..., -4.0114e-02,6.7487e-03, 1.0511e-03],[-2.4961e-02, 1.4747e-03, -5.4
所以如果我們用
torch.save(model.state_dict(), PATH)
保存我們的模型,我們可以用 Python 的 pathlib
模塊來測量它的大小。具體來說,Path(PATH).stat().st_size
返回文件大小,單位是字節。
讓我們將其集成到 PerformanceBenchmark
類中的一個 compute_size()
方法中,以自動化這個過程:
import torch
from pathlib import Path
def compute_size(self):state_dict = self.pipeline.model.state_dict()tmp_path = Path("model.pt")torch.save(state_dict, tmp_path)# 計算大小,單位為兆字節size_mb = Path(tmp_path).stat().st_size / (1024 * 1024)# 刪除臨時文件tmp_path.unlink()print(f"模型大小 (MB) - {size_mb:.2f}")return {"size_mb": size_mb}
PerformanceBenchmark.compute_size = compute_size
為了完成我們的基準測試,我們將測量推理延遲——模型處理單個輸入并返回預測意圖所需的時間。這為我們提供了一個關于真實世界響應性的估計,尤其是在需要實時預測的生產系統中尤為重要。
在這種情況下,延遲包括管道中的所有處理步驟,包括分詞和模型推理。雖然分詞速度極快(通常比推理快約 1000 倍),但它仍然是端到端過程的一部分,所以我們為了完整性而將其包含在內。
為了準確測量執行時間,我們將使用 Python 的 time.perf_counter()
,它提供高分辨率計時,比 time.time()
更適合性能基準測試。
我們可以通過傳遞測試查詢并計算開始和結束之間的時間差(以毫秒為單位)來用 perf_counter 對管道進行計時:
from time import perf_counter
for _ in range(3):start_time = perf_counter()_ = pipe(query)latency = perf_counter() - start_timeprint(f"延遲 (ms) - {1000 * latency:.3f}")
延遲 (ms) - 64.923
延遲 (ms) - 47.636
延遲 (ms) - 47.344
延遲在不同運行之間可能會有很大差異,特別是對于小輸入或在系統負載不一致的情況下。對管道進行單次傳遞的計時通常會因為背景進程、CPU 節流或即時編譯(JIT)效應而產生噪聲測量結果。
為了緩解這種情況并獲得更可靠的延遲估計,我們采取以下方法:
- 預熱 CPU:運行幾次初始推理以穩定運行時環境。
- 重復測量:對許多樣本進行推理,以收集延遲的分布。
- 報告均值和標準差:這些統計值提供了典型延遲及其可變性的更穩健視圖。
以下是如何在 PerformanceBenchmark
類中實現此邏輯:
import numpy as np
def time_pipeline(self, query="我的賬戶的 PIN 碼是多少?"):latencies = []# 預熱for _ in range(10):_ = self.pipeline(query)# 定時運行for _ in range(100):start_time = perf_counter()_ = self.pipeline(query)latency = perf_counter() - start_timelatencies.append(latency)# 計算運行統計信息time_avg_ms = 1000 * np.mean(latencies)time_std_ms = 1000 * np.std(latencies)print(f"平均延遲 (ms) - {time_avg_ms:.2f} +\- {time_std_ms:.2f}")return {"time_avg_ms": time_avg_ms, "time_std_ms": time_std_ms}
PerformanceBenchmark.time_pipeline = time_pipeline
對基線模型進行基準測試
我們將把結果收集到 perf_metrics 字典中,以便跟蹤每個模型的性能:
pb = PerformanceBenchmark(pipe, clinc["test"])
perf_metrics = pb.run_benchmark()
模型大小 (MB) - 418.17
平均延遲 (ms) - 46.05 +\- 10.13
測試集上的準確率 - 0.867
擴展智能:知識蒸餾用于高效模型部署
知識蒸餾是一種通用方法,用于訓練一個較小的學生模型來模仿一個較慢、較大但性能更好的教師模型的行為。
知識蒸餾用于高效微調
知識蒸餾是一種強大的技術,用于監督學習的微調階段,其中較大的、經過良好訓練的“教師”模型將其學到的行為傳遞給較小的“學生”模型。目標不僅僅是復制性能——而是傳遞通常在真實標簽中看不見的細微、學到的見解。
🔢 蒸餾背后的數學機制
- 生成 logits:輸入序列 x 被傳遞給教師,生成原始預測分數:z(x)=[z1?(x),z2?(x),…,zN?(x)]
- 帶溫度縮放的 softmax:傳統 softmax:
然而,這通常會導致尖峰分布,幾乎沒有信息增益。
改進的 softmax 帶溫度 T:
? 更高的 T? 更柔和的分布 ? 更有信息量,關于類別關系和決策邊界
?? 損失函數:平衡準確性與見解
-
學生的軟預測:qi(x)
-
KL 散度損失(知識蒸餾損失):
-
因子 T2 對梯度幅度進行歸一化。
-
總學生損失:
🧠 推理階段
在推理時,溫度 T 重置為 1,以恢復標準的預測置信度。
在預訓練期間進行知識蒸餾:構建更智能、更小的模型
盡管知識蒸餾通常用于微調,但它在預訓練期間同樣有效——允許創建更緊湊、通用的模型,這些模型更快且更高效。
預訓練中的工作原理
- 一個大型預訓練教師(例如 BERT)將其對掩碼語言建模(MLM)的理解傳遞給一個較小的學生。
- 學生不僅從原始的 MLM 目標中學習,還從教師的行為模式和表示中學習。
DistilBERT 損失函數
在DistilBERT架構中,總損失結合了三個組成部分:
實際應用
由于我們已經有一個微調過的 BERT-base 模型,我們現在可以:
- 將其用作教師來指導一個較小的學生模型。
- 實現一個自定義的Trainer,它整合了交叉熵和蒸餾損失。
這種方法不僅加快了推理時間,還減少了資源使用——而沒有過多地犧牲性能。
在 PyTorch 中構建知識蒸餾 Trainer
為了在微調設置中實現知識蒸餾,我們擴展了 Hugging Face Trainer
類,添加了允許學生模型從預訓練的教師模型學習的額外組件。
要添加的關鍵組件
- 超參數:
alpha (α)
:平衡交叉熵和蒸餾損失(默認 = 0.5)。temperature (T)
:軟化 logits 以獲得更平滑的概率分布(默認 = 2.0)。
2.教師模型:
- 一個微調過的 BERT-base模型作為教師,學生將從中學習。
3.自定義損失函數:
- 結合交叉熵損失(針對真實標簽)與KL 散度(模仿教師輸出)。
逐步代碼實現
1. 自定義訓練參數
from transformers import TrainingArguments
class DistillationTrainingArguments(TrainingArguments):def __init__(self, *args, alpha=0.5, temperature=2.0, **kwargs):super().__init__(*args, **kwargs)self.alpha = alphaself.temperature = temperature
2. 帶有蒸餾邏輯的自定義 Trainer
import torch.nn as nn
import torch.nn.functional as F
from transformers import Trainerclass DistillationTrainer(Trainer):def __init__(self, *args, teacher_model=None, **kwargs):super().__init__(*args, **kwargs)self.teacher_model = teacher_modeldef compute_loss(self, model, inputs):outputs_stu = model(**inputs)loss_ce = outputs_stu.losslogits_stu = outputs_stu.logits# 教師前向傳播(不計算梯度)with torch.no_grad():outputs_tea = self.teacher_model(**inputs)logits_tea = outputs_tea.logits# 計算基于 KL 散度的蒸餾損失loss_fct = nn.KLDivLoss(reduction="batchmean")loss_kd = self.args.temperature ** 2 * loss_fct(F.log_softmax(logits_stu / self.args.temperature, dim=-1),F.softmax(logits_tea / self.args.temperature, dim=-1))# 損失加權求和return self.args.alpha * loss_ce + (1. - self.args.alpha) * loss_kd
幕后工作原理
- 教師預測:不計算梯度;它是一個固定的模型。
- 軟 logits:學生 logits 通過
log_softmax
,教師 logits 通過softmax
。 - KL 散度:衡量學生模仿教師軟化預測的接近程度。
- 損失混合:最終損失 =
α * 交叉熵 + (1 - α) * 蒸餾損失
。
選擇一個好的學生初始化
首先,我們需要對查詢進行分詞和編碼,所以讓我們實例化 DistilBERT 的分詞器并創建一個簡單的函數來處理預處理:
student_ckpt = "distilbert-base-uncased"
student_tokenizer = AutoTokenizer.from_pretrained(student_ckpt)
def tokenize_text(batch, tokenizer):return tokenizer(batch["text"], truncation=True)
clinc_enc = clinc.map(tokenize_text, batched=True, remove_columns=["text"],fn_kwargs={"tokenizer": student_tokenizer})
clinc_enc.rename_column_("intent", "labels")
在這里,我們移除了 text 列,因為我們不再需要它,我們還使用 fn_kwargs 參數指定了 tokenize_text 函數中應該使用的分詞器。我們還將 intent 列重命名為 labels,以便它可以被訓練器自動檢測。現在我們的文本已經處理好了,接下來要做的是實例化 DistilBERT 進行微調。由于我們將多次運行訓練器,我們將使用一個函數來初始化每次運行的模型:
import torch
from transformers import AutoConfig
num_labels = intents.num_classes
id2label = bert_model.config.id2label
label2id = bert_model.config.label2id
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
student_config = (AutoConfig.from_pretrained(student_ckpt, num_labels=num_labels,id2label=id2label, label2id=label2id))
def student_init():return (AutoModelForSequenceClassification.from_pretrained(student_ckpt, config=student_config).to(device))
我們需要定義在訓練期間跟蹤的指標,
def compute_metrics(pred):predictions, labels = predpredictions = np.argmax(predictions, axis=1)return accuracy_score.compute(predictions=predictions, references=labels)
最后,我們只需要定義訓練參數。為了熱身,我們將 α 設置為 1,看看 DistilBERT 在沒有任何教師信號的情況下表現如何:
batch_size = 48
student_training_args = DistillationTrainingArguments(output_dir="checkpoints", evaluation_strategy = "epoch", num_train_epochs=5,learning_rate=2e-5, per_device_train_batch_size=batch_size,per_device_eval_batch_size=batch_size, alpha=1, weight_decay=0.01)
接下來我們加載教師模型,實例化訓練器并開始微調:
teacher_checkpoint = "lewtun/bert-base-uncased-finetuned-clinc"
teacher_model = (AutoModelForSequenceClassification.from_pretrained(teacher_checkpoint, num_labels=num_labels).to(device))
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();
將其包裝在 TextClassificationPipeline 中并通過我們的性能基準進行測試:
pipe = TextClassificationPipeline(model=distil_trainer.model.to("cpu"), tokenizer=distil_trainer.tokenizer)
optim_type = "DistilBERT"
pb = PerformanceBenchmark(pipe, clinc["test"], optim_type=optim_type)
perf_metrics.update(pb.run_benchmark())模型大小 (MB) - 255.89
平均延遲 (ms) - 24.13 +\- 10.06
測試集上的準確率 - 0.856
創建一個散點圖,將準確率與延遲進行對比,每個點的半徑對應模型的大小。
import pandas as pd
def plot_metrics(perf_metrics, current_optim_type):df = pd.DataFrame.from_dict(perf_metrics, orient='index')for idx in df.index:df_opt = df.loc[idx]if idx == current_optim_type:plt.scatter(df_opt["time_avg_ms"], df_opt["accuracy"] * 100,alpha=0.5, s=df_opt["size_mb"], label=idx,marker='$\u25CC')else:plt.scatter(df_opt["time_avg_ms"], df_opt["accuracy"] * 100,s=df_opt["size_mb"], label=idx, alpha=0.5)legend = plt.legend(bbox_to_anchor=(1,1))for handle in legend.legendHandles:handle.set_sizes([20])plt.ylim(80,90)plt.xlim(5, 53)plt.ylabel("準確率 (%)")plt.xlabel("平均延遲 (ms)")plt.show()
plot_metrics(perf_metrics, optim_type)
使用 Optuna 調整蒸餾超參數
Optuna 將超參數調整視為一個優化問題。它定義了一個目標函數,然后運行多次試驗以最小化或最大化它。
Rosenbrock 的香蕉函數:
優化中的一個經典基準:
- 全局最小值在:(x,y)=(1,1)
- 因其彎曲的香蕉形狀輪廓而得名
- 理論上簡單,但收斂到真實最小值具有挑戰性
現在,讓我們用類似的方法來優化 Hugging Face Trainer
中的知識蒸餾參數。
定義超參數空間
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)}
運行超參數搜索
best_run = distil_trainer.hyperparameter_search(n_trials=20,direction="maximize",hp_space=hp_space
)
direction="maximize"
告訴 Optuna 尋找更高的準確率。best_run
包含最佳試驗的配置和性能。
樣本輸出
print(best_run)
# BestRun(run_id='4', objective=3080.87,
# hyperparameters={'num_train_epochs': 8, 'alpha': 0.31, 'temperature': 16})
💡 洞見:所選的 α = 0.31 表明大部分學習信號來自知識蒸餾,而不是真實標簽。
應用最佳超參數并重新訓練
for k, v in best_run.hyperparameters.items():setattr(distil_trainer.args, k, v)distil_trainer.train()
保存模型以供日后使用:
distil_trainer.save_model("models/distilbert-base-uncased-distilled-clinc")
對蒸餾后的模型進行基準測試
創建一個管道并重新進行基準測試,看看我們在測試集上的表現如何:
pipe = TextClassificationPipeline(model=distil_trainer.model.to("cpu"), tokenizer=distil_trainer.tokenizer)
optim_type = "蒸餾"
pb = PerformanceBenchmark(pipe, clinc["test"], optim_type=optim_type)
perf_metrics.update(pb.run_benchmark())
模型大小 (MB) - 255.89
平均延遲 (ms) - 24.58 +\- 7.66
測試集上的準確率 - 0.871
用量化加速 Transformer
雖然知識蒸餾通過訓練一個較小的學生模型來減小模型大小,但量化通過降低計算精度——通常從 32 位浮點數(FP32)降低到 8 位整數(INT8)來提高效率。這可以帶來:
- 更小的模型大小
- 更快的推理速度
- 最小的準確率損失
可視化權重分布以進行量化
Transformer 權重通常位于一個狹窄的范圍內,使其非常適合 INT8 量化。
import matplotlib.pyplot as plt
weights = bert_model.state_dict()["bert.encoder.layer.0.attention.output.dense.weight"]
plt.hist(weights.flatten().numpy(), bins=250, range=(-0.3, 0.3));
如果大多數值位于 [?0.1, 0.1] 范圍內,我們可以安全地將它們量化為 INT8(?128 到 127),而幾乎沒有損失。
手動量化示例
步驟 1:計算比例因子和零點
zero_point = 0
scale = (weights.max() - weights.min()) / (127 - (-128))
步驟 2 :量化張量
(weights / scale + zero_point).clamp(-128, 127).round().char()
[[ 2, -1, 1, ..., -2, -6, 9],[ 7, 2, -4, ..., -3, 5, -3],[-15, -8, 5, ..., 3, 0, -2],...,[ 11, -1, 12, ..., -2, 0, -3],[ -2, -6, -13, ..., 11, -3, -10],[-12, 5, -3, ..., 7, -3, -1]], dtype=torch.int8)
步驟 3:使用 PyTorch 的 API
from torch import quantize_per_tensor
quantized_weights = quantize_per_tensor(weights, scale, zero_point, torch.qint8)
quantized_weights.int_repr()
([[ 2, -1, 1, ..., -2, -6, 9],[ 7, 2, -4, ..., -3, 5, -3],[-15, -8, 5, ..., 3, 0, -2],...,[ 11, -1, 12, ..., -2, 0, -3],[ -2, -6, -13, ..., 11, -3, -10],[-12, 5, -3, ..., 7, -3, -1]], dtype=torch.int8)
如果我們對這個張量進行反量化,我們可以可視化頻率分布,看看四舍五入對原始值的影響:
from mpl_toolkits.axes_grid1.inset_locator import zoomed_inset_axes,mark_inset
# 創建直方圖
fig, ax = plt.subplots()
ax.hist(quantized_weights.dequantize().flatten().numpy(),bins=250, range=(-0.3,0.3));
# 創建放大插入圖
axins = zoomed_inset_axes(ax, 5, loc='upper right')
axins.hist(quantized_weights.dequantize().flatten().numpy(),bins=250, range=(-0.3,0.3));
x1, x2, y1, y2 = 0.05, 0.1, 500, 2500
axins.set_xlim(x1, x2)
axins.set_ylim(y1, y2)
axins.axes.xaxis.set_visible(False)
axins.axes.yaxis.set_visible(False)
mark_inset(ax, axins, loc1=2, loc2=4, fc="none", ec="0.5")
plt.show()
這非常清楚地顯示了由于只精確映射一些權重值并對其余值進行四舍五入而引起的離散化。為了完善我們的小分析,讓我們比較一下計算兩個權重張量乘法所需的時間,一個使用 FP32 值,另一個使用 INT8 值。對于 FP32 張量,我們可以使用 PyTorch 的便捷 @ 運算符進行乘法:
%%timeit
weights @ weights
對于量化張量,我們需要 QFunctional 包裝器類,以便我們可以使用特殊的 torch.qint8 數據類型進行操作:
from torch.nn.quantized import QFunctional
q_fn = QFunctional()
這個類支持各種基本操作,比如加法,在我們的情況下,我們可以這樣對量化張量的乘法進行計時:
%%timeit
q_fn.mul(quantized_weights, quantized_weights)
107 μs ± 7.87 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
內存比較
import sys
sys.getsizeof(weights.storage()) / sys.getsizeof(quantized_weights.storage())
# 約小 4 倍
使用 PyTorch 量化 Transformer
from torch.quantization import quantize_dynamic
from transformers import AutoModelForSequenceClassification, AutoTokenizermodel_ckpt = "models/distilbert-base-uncased-distilled-clinc"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)model = AutoModelForSequenceClassification.from_pretrained(model_ckpt).to("cpu")model_quantized = quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
這行代碼:
- 量化所有
nn.Linear
層。 - 使用 INT8 算術進行更快的推理。
- 幾乎保持了相同的準確率。
對量化模型的性能進行基準測試
我們的模型已經成功量化,現在是時候測試它的性能了。我們將運行一個基準測試來評估它的速度和內存效率——這對于在資源受限的環境中部署至關重要。
以下是設置和執行基準測試的方式:
pipe = TextClassificationPipeline(model=model_quantized, tokenizer=tokenizer)
optim_type = "蒸餾 + 量化"
pb = PerformanceBenchmark(pipe, clinc["test"], optim_type=optim_type)
perf_metrics.update(pb.run_benchmark())
plot_metrics(perf_metrics, optim_type)
使用 ONNX 和 ONNX 運行時優化推理
我們的蒸餾模型已經經過優化和量化,現在是時候使用 ONNX 框架進一步突破極限了——這是一個強大的平臺,用于深度學習模型的互操作性和高性能推理。
ONNX(Open Neural Network Exchange)是一個開放標準,定義了:
- 跨框架的通用操作符集
- 統一的文件格式用于模型導出/導入
- 神經網絡的計算圖表示
得益于 ONNX,你可以輕松地導出 PyTorch 模型并將其導入到 TensorFlow 中——反之亦然——從而實現在不同生態系統中的靈活部署。
設置 OpenMP 環境變量以供 ONNX 使用:
from psutil import cpu_count
%env OMP_NUM_THREADS={cpu_count()}
%env OMP_WAIT_POLICY=ACTIVE
env: OMP_NUM_THREADS=8
env: OMP_WAIT_POLICY=ACTIVE
將我們的蒸餾模型轉換為 ONNX 格式:
from transformers.convert_graph_to_onnx import convert
onnx_model_path = Path("onnx/model.onnx")
convert(framework="pt", model=model_ckpt, tokenizer=tokenizer,output=onnx_model_path, opset=12, pipeline_name="sentiment-analysis")
ONNX opset version set to: 12
Loading pipeline (model: models/distilbert-base-uncased-distilled-clinc,> tokenizer: PreTrainedTokenizerFast(name_or_path='models/distilbert-base-> uncased-distilled-clinc', vocab_size=30522, model_max_len=512, is_fast=True,> padding_side='right', special_tokens={'unk_token': '[UNK]', 'sep_token':> '[SEP]', 'pad_token': '[PAD]', 'cls_token': '[CLS]', 'mask_token':> '[MASK]'}))
Creating folder onnx
Using framework PyTorch: 1.5.0
Found input input_ids with shape: {0: 'batch', 1: 'sequence'}
Found input attention_mask with shape: {0: 'batch', 1: 'sequence'}
Found output output_0 with shape: {0: 'batch'}
Ensuring inputs are in correct order
head_mask is not present in the generated input list.
Generated inputs order: ['input_ids', 'attention_mask']
ONNX 使用操作符集將不可變的操作符規范分組在一起,因此 opset=12 對應于 ONNX 庫的一個特定版本。現在我們已經保存了模型。
我們需要創建一個推理會話來將輸入傳遞給模型:
om onnxruntime import (GraphOptimizationLevel, InferenceSession,SessionOptions)
def create_model_for_provider(model_path, provider="CPUExecutionProvider"):options = SessionOptions()options.intra_op_num_threads = 1options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALLsession = InferenceSession(str(model_path), options, providers=[provider])session.disable_fallback()return session
onnx_model = create_model_for_provider(onnx_model_path)
用測試集中的一個示例進行測試。由于轉換函數的輸出告訴我們 ONNX 只期望 input_ids
和 attention_mask
作為輸入,因此我們需要從樣本中刪除標簽列:
inputs = clinc_enc["test"][:1]
del inputs["labels"]
logits_onnx = onnx_model.run(None, inputs)[0]
logits_onnx.shape##(1, 151)
通過取 argmax 獲取預測標簽:
np.argmax(logits_onnx)## 添加真實標簽clinc_enc["test"][0]["labels"]
我們將創建自己的類來模擬核心行為:
from scipy.special import softmax
class OnnxPipeline:def __init__(self, model, tokenizer):self.model = modelself.tokenizer = tokenizerdef __call__(self, query):model_inputs = self.tokenizer(query, return_tensors="pt")inputs_onnx = {k: v.cpu().detach().numpy()for k, v in model_inputs.items()}logits = self.model.run(None, inputs_onnx)[0][0, :]probs = softmax(logits)pred_idx = np.argmax(probs).item()return [{"label": intents.int2str(pred_idx), "score": probs[pred_idx]}]
然后我們可以在簡單的查詢上測試這個,看看我們是否能夠恢復 car_rental
意圖:
pipe = OnnxPipeline(onnx_model, tokenizer)
pipe(query)
[{'label': 'car_rental', 'score': 0.8440852}]
高效地對 ONNX 模型進行基準測試
現在我們的 ONNX 管道已經正常工作,下一步是對它的性能進行基準測試。為此,我們將擴展我們現有的 PerformanceBenchmark
類。由于 ONNX 模型是一個 InferenceSession
實例(而不是 PyTorch 的 nn.Module
),它沒有像 state_dict
這樣的屬性,因此無法使用 torch.save()
來計算大小。
🔧 為了解決這個問題,我們將只覆蓋 compute_size()
方法,同時重用現有的 compute_accuracy()
和 time_pipeline()
的實現。
以下是一種簡潔的方式來處理 ONNX 模型的大小計算:
lass OnnxPerformanceBenchmark(PerformanceBenchmark):def __init__(self, *args, model_path, **kwargs):super().__init__(*args, **kwargs)self.model_path = model_pathdef compute_size(self):size_mb = Path(self.model_path).stat().st_size / (1024 * 1024)print(f"模型大小 (MB) - {size_mb:.2f}")return {"size_mb": size_mb}
使用我們新的基準測試工具,讓我們看看將蒸餾模型轉換為 ONNX 格式后的表現如何:
optim_type = "蒸餾 + ORT"
pb = OnnxPerformanceBenchmark(pipe, clinc["test"], optim_type,model_path="onnx/model.onnx")
perf_metrics.update(pb.run_benchmark())# 模型大小 (MB) - 255.89
# 平均延遲 (ms) - 10.54 +\- 2.20
# 測試集上的準確率 - 0.871plot_metrics(perf_metrics, optim_type)
使用 ONNX 運行時優化 Transformer 推理
我們已經看到,當將蒸餾 Transformer 模型轉換為 ONNX 格式時,ONNX 運行時(ORT)已經提供了相當不錯的性能提升。但我們還可以更進一步,通過應用 ORT 的優化工具包中的Transformer 特定圖優化。
Transformer 特定優化
對于像 DistilBERT 這樣的 Transformer 架構,ONNX 運行時工具提供了針對類型為 bert
的模型的高級優化。首先,我們使用 BertOptimizationOptions
定義一組模型特定的優化選項:
from onnxruntime_tools.transformers.onnx_model_bert import BertOptimizationOptionsmodel_type = "bert"
opt_options = BertOptimizationOptions(model_type)
opt_options.enable_embed_layer_norm = False # 改善模型大小壓縮
禁用嵌入層歸一化融合在某些情況下可以實現更好的壓縮效果。
接下來,我們運行優化過程:
from onnxruntime_tools import optimizeropt_model = optimizer.optimize_model("onnx/model.onnx",model_type=model_type,num_heads=12,hidden_size=768,optimization_options=opt_options
)opt_model.save_model_to_file("onnx/model.opt.onnx")
我們提供了 DistilBERT 模型的注意力頭數和隱藏層大小。優化完成后,我們可以運行性能基準測試:
onnx_model_opt = create_model_for_provider("onnx/model.opt.onnx")
pipe = OnnxPipeline(onnx_model_opt, tokenizer)
optim_type = "蒸餾 + ORT (優化)"pb = OnnxPerformanceBenchmark(pipe, clinc["test"], optim_type, model_path="onnx/model.opt.onnx")
perf_metrics.update(pb.run_benchmark())# 輸出# 模型大小 (MB) - 255.86# 平均延遲 (ms) - 11.22 ± 3.52# 測試集上的準確率 - 0.871plot_metrics(perf_metrics, optim_type)
🔍 洞見:我們最初的 ONNX 優化已經接近最優——這個針對 BERT 的特定優化并沒有在大小或速度上帶來重大改進。
加入量化
為了進一步減小大小和延遲,我們使用 ONNX 運行時的量化工具應用動態量化。與 PyTorch 主要量化 nn.Linear
層不同,ORT 還可以量化嵌入層,從而獲得更好的結果。
from onnxruntime.quantization import quantize_dynamic, QuantTypemodel_input = "onnx/model.onnx"
model_output = "onnx/model.quant.onnx"quantize_dynamic(model_input, model_output, weight_type=QuantType.QInt8)
現在,讓我們對量化后的 ONNX 模型進行基準測試:
onnx_quantized_model = create_model_for_provider(model_output)
pipe = OnnxPipeline(onnx_quantized_model, tokenizer)
optim_type = "蒸餾 + ORT (量化)"pb = OnnxPerformanceBenchmark(pipe, clinc["test"], optim_type, model_path=model_output)
perf_metrics.update(pb.run_benchmark())# 輸出# 模型大小 (MB) - 185.71# 平均延遲 (ms) - 6.95 ± 4.75# 測試集上的準確率 - 0.875plot_metrics(perf_metrics, optim_type)
結果:ORT 量化將大小和延遲都幾乎減少了 50%,與 PyTorch 量化相比。總體而言,這帶來了令人印象深刻的 7 倍加速,與原始 BERT 基線相比,準確率幾乎沒有損失。