模型壓縮:蒸餾、剪枝、量化 - 把模型塞進邊緣攝影機的三條路

模型壓縮:蒸餾、剪枝、量化 - 把模型塞進邊緣攝影機的三條路

邊緣攝影機的盒子裡沒有 GPU,記憶體以 KB 計價,雲端訓好的模型根本塞不進去。本篇實測模型壓縮三招:知識蒸餾(讓小模型抄大模型的軟答案)、剪枝(拔掉九成權重)、量化(FP32 變 INT8)。結果誠實到反直覺:蒸餾在小資料上五組配方全數打平沒賺、小模型量化直接崩盤;真正的贏家是「壓縮大模型」——剪 90% 還有 0.833,把同尺寸的小模型(0.583)遠遠甩開。


核心貢獻者

Geoffrey Hinton、Oriol Vinyals 與 Jeff Dean(Google)於 2015 年發表《Distilling the Knowledge in a Neural Network》,提出知識蒸餾:大模型輸出的「軟機率」藏著類別之間的相似結構(暗知識),小模型抄這份答案比抄硬標籤學得更多。同年,韓松 (Song Han) 等人(Stanford)發表 Deep Compression,用「剪枝 + 量化 + 霍夫曼編碼」把 AlexNet 壓小 35 倍且不掉精度,拿下 ICLR 2016 最佳論文——他後來在 MIT 開創的 TinyML 實驗室,就是「把 AI 塞進微控制器」這條路線的大本營。


為什麼要壓縮?

Day 44 到 48 一路把模型越做越強:ViT、CLIP、SAM——但它們全都住在雲端或工作站。回到工安落地的最後一哩:廠區門口那台邊緣攝影機,盒子裡沒有 GPU、快閃記憶體以 KB 到 MB 計價、還要求毫秒級反應。雲端訓好的模型直接搬過去,第一步就卡死:塞不進去

把模型變小,業界有三條標準路,而且可以疊加:

Loading Diagram...

本篇把三招全部實測一遍,包含它們疊起來會發生什麼——結果比教科書誠實得多。


1. 資料集來源

資料集:Olivetti Faces (AT&T)

備註:Day 45 的老朋友,40 人、每人 10 張、64×64 灰階。這次任務改成 40 人分類(每人 7 張訓練、3 張測試),當作「門禁攝影機要離線辨識員工」的代理情境。真實邊緣場景跑的是 YOLO 這類偵測器(本篇環境跑不動),但三招壓縮的數學與取捨,和 YOLO 上一模一樣

實驗設計:一胖一瘦兩個模型

  • 老師 (Teacher):三層 CNN,98,280 參數,FP32 共 384 KB——假想它是「塞不進盒子」的雲端模型。
  • 學生 (Student):同構但通道數砍到 (8, 16, 32),7,320 參數,29 KB,小 13.4 倍
  • 三招各自上陣:學生抄老師(蒸餾)、老師拔權重(剪枝)、兩者都轉 INT8(量化),統一在 120 張測試集上驗收。

2. 原理

2.1 知識蒸餾:軟答案裡的暗知識

老師對一張人臉的輸出,不只是「這是 5 號員工」,還包含「他跟 12 號有點像、跟 30 號完全不像」——這份類別間的相似結構,Hinton 稱為暗知識 (Dark Knowledge),硬標籤裡完全沒有。用溫度 T 把 softmax 軟化,暗知識才會浮出來:

暗知識

左邊 T=1 幾乎是贏者全拿;右邊 T=4 之後,「哪些人長得像」的資訊整個攤開。學生的損失是兩份作業的加權:

L=αCE(y,ps)+(1α)T2KL ⁣(pt(T)ps(T))L = \alpha \cdot \text{CE}(y, p_s) + (1-\alpha) \cdot T^2 \cdot \text{KL}\!\big(p_t^{(T)} \,\|\, p_s^{(T)}\big)

  • 第一項抄硬標籤(標準答案),第二項抄老師的軟輸出(解題思路)。
  • T2T^2 是補償項:溫度拉高會讓梯度變小,乘回去才公平。

2.2 剪枝:神經網路的斷捨離

訓練完的網路裡,大量權重接近零、近乎沒在工作。幅度剪枝 (Magnitude Pruning) 就是全域排序所有權重的絕對值,最小的那批直接歸零。關鍵在剪完之後:

  • 直接剪 = 腦葉切除:剪 70% 不微調,準確率從 0.925 崩到 0.067。
  • 剪完微調 = 復健:讓存活的權重重新分工,幾十個 epoch 就能大幅回血。

2.3 量化:FP32 → INT8,大小除以 4

每個權重從 32-bit 浮點數改用 8-bit 整數表示:記下每組權重的縮放係數 scale,權重四捨五入到 254 個整數格點上。本篇採用逐通道對稱量化(每個輸出通道自己一個 scale),這是業界部署的標準做法——模型直接小 4 倍,而且整數運算在邊緣晶片上通常更快、更省電。

2.4 三招對照

面向蒸餾剪枝量化
改變什麼換一個小架構拔掉部分權重降低數值精度
需要重訓要(學生全程重訓)微調即可免訓練(PTQ)
省多少看學生大小(本篇 13 倍)看稀疏度(50~95%)固定 4 倍
隱藏成本要調 α 與 T非結構化剪枝不加速小模型可能崩

3. 實戰

Python 程式碼實作

# 蒸餾:硬標籤 + 老師軟輸出的加權(Hinton 2015)
kd = F.kl_div(F.log_softmax(student_logits / T, 1),
              F.softmax(teacher_logits / T, 1),
              reduction='batchmean') * T * T
loss = alpha * F.cross_entropy(student_logits, y) + (1 - alpha) * kd

# 剪枝:全域幅度排序,最小的一批歸零;微調時每步都把遮罩壓回去
def global_mask(model, sparsity):
    ws = [(n, p) for n, p in model.named_parameters() if p.dim() > 1]
    allw = torch.cat([p.abs().flatten() for _, p in ws])
    thr = torch.quantile(allw, sparsity)
    return {n: (p.abs() > thr).float() for n, p in ws}

# 量化:逐通道對稱 INT8(業界標準),權重四捨五入到整數格點
for p in model.parameters():
    if p.dim() > 1:
        flat = p.reshape(p.shape[0], -1)
        scale = (flat.abs().max(1).values / 127.0).clamp(min=1e-12)
        p.copy_(((flat / scale[:, None]).round().clamp(-127, 127)
                 * scale[:, None]).reshape(p.shape))

程式碼重點:

  • 蒸餾的 T * T 不能漏,漏了軟目標的梯度會被溫度稀釋。
  • 剪枝微調時,每個 optimizer step 後都要重新套遮罩,不然被剪掉的權重會偷偷復活。
  • 量化用「假量化」模擬:權重真的落在 INT8 格點上,但用浮點儲存做推論——這是評估精度損失的標準手法;真上邊緣還需要整數推論引擎(TFLite、ONNX Runtime)。

4. 模型評估

總成績單(120 張測試集)

版本準確率大小 (KB)CPU 延遲 (ms/張)
Teacher FP32(基準)0.9253842.0
Teacher INT80.933962.2
Teacher 剪 90% + 微調0.833382.0
Student 從頭訓練0.583291.2
Student 蒸餾(最佳配方)0.583291.2
Student 蒸餾 + INT80.39271.1

剪枝:拔掉九成,還能站著

剪枝掃描
  • 橘線(不微調):剪 70% 就腦死(0.067)——權重之間有精密分工,亂拔當然崩。
  • 綠線(微調 30 epochs):剪 50%、70% 完全無損(0.933,甚至比基準高)、剪 90% 仍有 0.833、剪到 95% 才真正斷崖(0.492)。網路過參數化的程度,比直覺誇張得多
  • 誠實註記:此為非結構化剪枝,省的是儲存(38 KB),延遲不變(表中同為 2.0 ms)——一般硬體不會因為權重是零就跳過計算;要換取速度得用結構化剪枝(整條通道拔掉)。

部署地圖:哪個點該上車?

部署地圖

右上是甜蜜區:Teacher INT8(96 KB,0.933)幾乎免費——大小除以 4,精度分毫未損(+0.008 在 120 張測試集上是 1 張圖的雜訊範圍)。Teacher 剪 90%(38 KB,0.833)是第二志願。而左下角的學生們,便宜但傷痕累累。

誠實的打臉(一):蒸餾五連敗

這篇本來的主角是蒸餾,結果五組配方全軍覆沒:

配方 (α, T)學生準確率vs 從頭訓練 (0.583)
α=0.3, T=40.450
α=0.5, T=40.583平手
α=0.5, T=20.558
α=0.7, T=20.583平手
α=0.0, T=4(純蒸餾)0.375慘敗

診斷:小資料是蒸餾的死穴。每人只有 7 張訓練照,硬標籤已經把資訊榨乾,老師的暗知識沒有額外的資料讓學生消化;Hinton 論文的紅利來自海量資料(甚至無標註遷移集)+ 集成老師,那才是蒸餾的主場。玩具規模硬上,只是讓學生同時抄兩份互相干擾的筆記。

誠實的打臉(二):小模型禁不起量化

同樣的 INT8:老師 0.925 → 0.933(無感),學生 0.583 → 0.392(崩盤)。原因是冗餘度:過參數化的大網路,四捨五入的誤差有大把冗餘權重吸收;已經瘦到 7,320 參數的學生,每個權重都身兼要職,一動就痛。「壓縮技巧疊好疊滿」是危險的直覺——越小的模型,對每一次再壓縮越脆弱。

工安視角的解讀

  • 邊緣選型法則:與其為邊緣攝影機「從頭訓練一個小模型」,不如「訓練大模型再壓縮」——本實驗中 38 KB 的剪枝老師(0.833)把 29 KB 的學生(0.583)徹底擊敗,這正是 Deep Compression 的論文主旨,在我們自己的數據上重演了一次。
  • 落地順序:先 INT8 量化(幾乎免費、直接省 4 倍)→ 不夠再剪枝 + 微調 → 蒸餾留給「有大量資料」的場景。
  • 量化必須驗收:上線前一定要用測試集重新評估,尤其模型本來就小的時候——0.583 變 0.392 這種事,不會在說明書上警告你。
  • 真實世界的對應:YOLO 家族的邊緣部署(YOLOv8n + INT8 TensorRT/TFLite)走的就是同一套流程;工安攝影機常見的「NPU 盒子」(如 Hailo、瑞芯微)吃的正是量化後的模型。

5. 總結

我們學習了模型壓縮三招:

  • 蒸餾:小模型抄大模型的軟答案(暗知識),但主場在大資料——小資料上五組配方全數打平或更差。
  • 剪枝:微調是靈魂——剪 70% 無損、剪 90% 仍有 0.833;但非結構化剪枝只省空間不省時間。
  • 量化:INT8 幾乎免費(大模型),但小模型會崩(0.583 → 0.392)——壓縮不能疊好疊滿。
  • 最大教訓:「訓大再壓」勝過「直接訓小」——38 KB 的剪枝老師輾壓 29 KB 的學生,Deep Compression 誠不我欺。

下一篇處理一個更根本的部署問題:模型再小再快,它答錯時毫無自覺——門禁放錯人、裂縫漏報,信心分數照樣 0.98。能不能讓模型誠實說「這題我不確定,請人來看」?不確定性估計與 Conformal Prediction——給預測加上統計保證的安全網。