推論引擎與加速器：ONNX、TensorRT 與硬體堆疊

從框架到引擎：部署用的運算圖

當你在 PyTorch 或 TensorFlow 裡訓練模型時，你跑的其實就是一般的 Python。每一行——這裡一個卷積、那裡一個 ReLU——都會立刻、一次一個運算地執行。這叫做即時（eager）執行，對研究來說非常好用：你可以在網路中途印出張量、設中斷點、兩次執行之間就改架構。但同樣這份彈性，正是讓模型難以部署、又慢又重的原因。每一次前向傳播都要付出 Python 直譯器的成本、付出它早已不需要的自動微分記帳成本，還要扛著一個為了訓練、而非為了服務而打造的龐大框架。

部署時我們做的是另一件事：把模型匯出（export）成一張凍結的運算圖（computational graph），再把這張圖交給精簡的推論引擎。運算圖其實就是把模型做的數學畫成的有向圖。節點（node）是運算子——`conv`、`relu`、`add`、`matmul`——邊（edge）則是在它們之間流動的張量（前一個運算子的輸出就是下一個的輸入）。數字本身完全沒變；我們只是把這串固定的步驟寫了下來，而不是每次推論都在 Python 裡重新決定一遍。

這張凍結的圖，就是所有部署工具的原料。這份指南會一路跟著它走完整個堆疊：先是 ONNX，一種標準檔案格式，讓圖不再被綁死在單一框架裡（由 ONNX Runtime 來執行）；接著是改寫圖、讓它跑更快的圖優化；再來是像 TensorRT 這類為某顆特定晶片編譯圖的廠商引擎；然後是把圖與硬體匹配的藝術；最後用剖析來確認加速是真的。整份指南想壓低的那個數字，就是 推論延遲——一次預測要花多久。

ONNX：模型的共通語言

在 PyTorch 裡訓練出的模型，預設就是個 PyTorch 的東西——要跑它你就得裝 PyTorch。ONNX（Open Neural Network Exchange，開放神經網路交換格式）打破了這種綁定。它是一種單一的標準檔案格式，用一種與產生它的框架無關的方式，描述那張運算圖——也就是第 1 節說的節點與邊。把你的 PyTorch 模型匯出成 `.onnx` 檔，你就能用 ONNX Runtime 在 Windows 筆電、Linux 伺服器、甚至 C++ 應用程式裡執行它，完全看不到 PyTorch 的影子。ONNX 把模型在哪裡訓練和它在哪裡執行這兩件事解耦了。

import torch

model.eval()  # turn off dropout / training-mode batchnorm behaviour
dummy = torch.randn(1, 3, 224, 224)  # one example with the input shape

torch.onnx.export(
    model, dummy, "model.onnx",
    input_names=["image"], output_names=["logits"],
    opset_version=17,                      # which operator dictionary to use
    dynamic_axes={"image": {0: "batch"}},  # let batch size vary at run time
)

注意 `opset_version` 這個參數。opset（運算子集合）是一本帶版本的字典，精確定義每個運算子的意思——`Conv` 吃哪些輸入、`Resize` 怎麼取整，諸如此類。ONNX 會持續演進，運算子會被新增、被修訂，每一次發布都被標上一個 opset 版本號。匯出端和執行端必須一致：如果你用 opset 17 匯出，但你的執行引擎只懂到 opset 13，或者你用了某個舊 opset 從未定義過的全新運算子，那匯出或載入就會直接失敗。大多數「我的 ONNX 匯出掛了」的問題都是 opset 不匹配——調高或調低 opset 版本，是第一個該試的。

ONNX Runtime 是真正執行這個檔案的引擎，它的關鍵法寶是執行提供者（execution provider，EP）——一組可插拔的後端，每一個都知道如何在某種特定硬體上跑運算子。CPU EP 在處理器上跑運算；CUDA EP 在 NVIDIA GPU 上跑；其他 硬體加速器 也各有對應的 EP。你給 ONNX Runtime 一份依優先序排列的 EP 清單，對每個運算子，它會挑出支援該運算、且優先序最高的提供者，遇到某個提供者跑不了時就往清單下面退。於是同一個 `model.onnx` 既能跑在只有 CPU 的機器上，也能跑在 GPU 伺服器上——你只要換提供者清單，而不必改模型。

import onnxruntime as ort

sess = ort.InferenceSession(
    "model.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)  # try the GPU first; fall back to CPU for any op CUDA can't run
out = sess.run(None, {"image": batch})[0]

圖優化：運算子融合與常數摺疊

接下來是讓人意外的部分：引擎可以在完全不改變它算出的任何一個數字的前提下，讓你的模型明顯變快。訣竅是把圖改寫成一個等價、但更便宜執行的形式。ONNX Runtime、TensorRT，以及每個認真的引擎，在載入你的模型時都會跑一整套這種圖優化。其中最重要的兩種，是運算子融合與常數摺疊。

運算子融合（operator fusion）把好幾個相鄰的運算子併成一個。一個經典樣式——在上面那個殘差塊裡就看得到——是 `Conv → BatchNorm → ReLU`。不融合的話，GPU 會算完整個 Conv 輸出，把整張張量寫出到慢速的全域記憶體，再全部讀回來做 BatchNorm，又寫一次，第三次再讀回來做 ReLU，然後再寫一次。融合之後，單一個核心（kernel）會在每一小塊資料還待在快速的晶片上暫存器裡時，就一口氣做完「卷積→正規化→截斷」，最後只寫出結果一次。回想第 1 篇的記憶體頻寬瓶頸：對這些層來說，晶片大部分時間花在搬資料、而不是在做乘法，所以把三趟往返記憶體砍成一趟，就算算術完全一樣，也能大幅壓低 推論延遲。

常數摺疊（constant folding）會在建置時（build time）先把任何輸入已知的計算算完，於是推論時就不必再跑。最招牌的例子，是把 BatchNorm 摺進前一個卷積裡。推論時的 BatchNorm 不過是一個固定的、每通道的縮放與平移（它的滑動平均與變異數都已凍結）。從代數上看，對卷積輸出做縮放與平移，等同於用一組稍微重新縮放過的卷積權重、再加上一個微調過的偏置——所以引擎在建圖時就把那組調整後的權重一次算好，BatchNorm 這個節點就整個消失了。輸出一模一樣，卻永遠少了一個運算子。

初學者常問：如果融合並沒有拿掉任何乘加運算，為什麼會變快？有兩個原因，都是關於額外開銷、而非 FLOPs。第一，記憶體流量——如上所述，更少的中間張量代表更少趟去慢速 DRAM，而許多視覺層是受限於記憶體、而非受限於計算的。第二，核心啟動開銷（kernel-launch overhead）——每個獨立的運算子都是一次 GPU 核心啟動，帶著固定的前置成本；發射一個融合核心、而不是三個，就省下兩次啟動。像 TensorRT 這類引擎，正是為了這些原因而極其積極地進行融合。

TensorRT 與廠商引擎

TensorRT 是 NVIDIA 的廠商引擎，它比 ONNX Runtime 又更進一步。ONNX Runtime 大致上是直譯你的圖——它對每個運算子都備有一個好用的核心，然後派工過去。TensorRT 則是把你的圖編譯成單一個與硬體綁定的引擎（engine）檔，為某一顆確切的 GPU 量身打造。可以想成：直譯式腳本，對比於針對你那顆確切 CPU、開滿所有最佳化旗標編譯出來的程式。

1. 自動調校核心（autotuning）：對每一層，TensorRT 都有許多候選實作（不同的分塊、演算法、資料佈局）。它會真的在你的 GPU 上把它們跑一遍計時，為那個確切的形狀與晶片留下最快的那個。
2. 逐層選擇精度：它可以讓某些層跑 FP16 或 INT8、其他層跑 FP32，逐層挑出在維持足夠準確度下最快的那個。
3. 積極融合：它把 conv+bias+activation 以及許多其他樣式融合得比一般引擎更徹底，因為它只需要在單一硬體目標上正確就好。

那個逐層精度，正是第 3 篇回來的地方。要讓某一層跑 INT8 推論，TensorRT 必須知道怎麼把每個張量的浮點數值範圍，映射到 256 個整數階上——它需要為每個激活張量找一個縮放係數（scale）。它透過一次校準（calibration）來學會這些係數：你餵給它幾百張有代表性的圖片，它記錄流經每一層的激活的真實範圍，再挑出資訊損失最小的縮放係數。這正是第 3 篇講的訓練後量化校準，只是現在由引擎建置器自動完成。沒有校準資料，就沒有安全的 INT8。

# Build an INT8 engine from an ONNX file, tuned for THIS machine's GPU.
trtexec \
  --onnx=model.onnx \
  --int8 \
  --calib=calib.cache \        # scales learned from representative images
  --saveEngine=model.plan       # the compiled, GPU-specific engine

代價是：一個 TensorRT 引擎和它建置時所在的硬體是綁死的。那個 `.plan` 檔綁定某個 GPU 架構、而且通常還綁定某個 TensorRT／驅動版本——把它搬到另一顆 加速器 上，它可能就拒絕載入，所以你得在每個目標上重新建置。建置本身也要花不少時間（從幾秒到好幾分鐘），就是因為那些自動調校。你在部署時付一次這個成本，換來那顆晶片上能拿到的最低延遲。對一個要回答上百萬次請求的服務來說，這幾乎總是值得的。

讓模型與加速器匹配

一張訓練好的圖，和一顆 加速器，並不保證合得來。同一個匯出的模型，在某顆晶片上可能快得驚人、在另一顆上卻慢得難堪——不是因為那顆晶片比較弱，而是因為模型用了那顆晶片並非為其打造、無法加速的運算子或形狀。要拿到好效能，有一部分是個協同設計（co-design）問題：你帶著目標硬體去塑造模型，就像你設計家具時會先想好它得穿過哪扇門。

加速器只會在它有專屬電路的運算上跑得快。GPU 的張量核心（tensor core）是專門做矩陣乘法的單元，跑卷積與矩陣乘法快得尖叫——但前提是形狀對得上它的分塊尺寸（例如通道數是 8 或 16 的倍數），而且精度是它支援的（FP16／INT8）。行動裝置上的 NPU 更挑：它只實作一份固定、有限的運算子菜單，菜單以外的任何東西，它的快速路徑就是跑不了。

1. 偏好硬體友善的運算子：標準卷積、ReLU、矩陣乘法到處都被加速；冷僻的激活與花俏的自訂層通常沒有。能換成被支援的等價物就換。
2. 讓形狀對齊硬體：把通道數保持在 8 或 16 的倍數，好讓張量核心填得滿；寧可補零（padding），也別留一個彆扭的 13 通道層。
3. 先查運算子支援表：在你決定一個架構之前，先讀你那版 TensorRT 或 NPU SDK 的支援運算子清單，確認 ONNX Runtime（或你的引擎）能把每一個運算都擺到加速器上。

把這點刻進腦裡：同一個 `.onnx` 檔，純粹因為這份匹配的好壞，就可能快、也可能慢。兩位工程師部署一模一樣的權重，卻看到 10 倍的延遲差距，只因為一位讓每個運算都留在加速器上、另一位在熱路徑裡留了一個退回。硬體感知的設計不是最後的拋光——它是你從第一張架構草圖就得帶著走的限制。

剖析已部署的引擎：閉合迴圈

你已經匯出、融合、編譯、也匹配好了。現在——就跟第 1 篇一樣——你要測量，因為在你親眼從數字裡看到之前，那些最佳化都還不算數。做這件事的工具，是逐層的 剖析（profiling）：不是在整個網路外面掛一支碼錶，而是替引擎裡每一個運算子計時，看毫秒究竟花到哪去了。

一份好的逐層剖析，能一次回答三個部署問題。融合有發生嗎？如果優化後的圖裡仍然分別列著 Conv、BatchNorm、ReLU，那你的融合就是悄悄沒觸發——也許是某個不支援的樣式擋住了它。現在的瓶頸是什麼？吃掉大部分時間的那一兩層，才是再下功夫會有回報的地方；其餘都是雜訊。有東西退回 CPU 嗎？跑在 CPU 提供者上的運算、或是資料被搬下加速器時突然出現的空檔，會把第 5 節那種無聲的退回給揪出來。

import onnxruntime as ort

opts = ort.SessionOptions()
opts.enable_profiling = True            # record per-op timings
sess = ort.InferenceSession(
    "model.onnx", opts,
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"])

for _ in range(50):                     # warm up, then run
    sess.run(None, {"image": batch})

prof = sess.end_profiling()             # JSON trace: per-op time + which EP ran it
# Open the trace, sort ops by duration, and look for CPUExecutionProvider rows.

最後，在做好的引擎上掃一遍批次大小（batch size）。回想第 1 篇的兩種速度：延遲 是一筆請求要等多久；吞吐量 是引擎每秒能穩定產出多少筆預測。較大的批次能讓加速器的計算單元更滿，所以吞吐量會往上爬——但每一筆請求現在都得等整批一起做完，所以延遲會上升。把引擎分別跑在 batch 1、2、4、8、16……並把兩者畫出來。對的批次大小，是仍能滿足你延遲預算的那個最大值；這就是在這個確切引擎上，把兩種速度互相權衡的那一個旋鈕。

1. 逐層剖析優化後的引擎。
2. 確認收穫：融合有觸發、也沒有意外的 CPU 退回。
3. 找出新的瓶頸運算子並處理它（換運算子、換形狀或換精度），或為你需要的吞吐量掃批次大小。
4. 重新匯出／重新建置、重新測量，反覆直到你達到延遲與吞吐量目標。

這個「部署—測量—調校」的迴圈，正是把一個只是跑得起來的模型，變成在它的目標上真正快的模型的關鍵。到目前為止，我們都假設有一台大方的伺服器、配著一顆大 GPU。第 5 篇會把這裡所有東西——匯出、融合、量化、運算子匹配、剖析——帶到最嚴苛的目標：邊緣（edge），那裡記憶體微小、電力稀缺，也沒有退回到一顆強悍 CPU 這條後路。紀律是一樣的，只是限制變得毫不留情。