ONNX Runtime 推理加速：8 個降低 Python 延遲的硬核技巧

在深度學(xué)習(xí)落地過程中，有一個常見的誤區(qū)：一旦推理速度不達(dá)標(biāo)，大家的第一反應(yīng)往往是拿著模型開到，比如：做剪枝、搞蒸餾、甚至犧牲精度換小模型。

實(shí)際上生產(chǎn)環(huán)境中的 Python 推理鏈路隱藏著巨大的“工程紅利”。很多時候你的模型本身并不慢，慢的是低效的數(shù)據(jù)搬運(yùn)、混亂的線程爭用以及不合理的 Runtime 默認(rèn)配置。在不改變模型精度的情況下，僅靠ONNX Runtime (ORT) 的工程特性，往往就能從現(xiàn)有技術(shù)棧中“摳”出驚人的性能提升。

以下是 8 個經(jīng)過實(shí)戰(zhàn)驗證的低延遲優(yōu)化策略，專治各種“莫名其妙的慢”。

1、 明確指定 Execution Provider 及其順序

ORT 會嚴(yán)格按照你傳入的 providers 列表順序進(jìn)行嘗試。把最快的放在第一位，并且盡量避免它靜默回退（Fallback）到 CPU。如果不顯式指定，ORT 有時候會“猶豫”，這都會消耗時間。

import onnxruntime as ort
providers = [
("TensorrtExecutionProvider", {"trt_fp16_enable": True}), # if supported
"CUDAExecutionProvider",
"CPUExecutionProvider",
]
sess = ort.InferenceSession("model.onnx", providers=providers)
print(sess.get_providers()) # verify what you actually got

Fallback 是有成本的，如果環(huán)境里有 TensorRT 就優(yōu)先用，沒有就降級到 CUDA，最后才是 CPU。把這個路徑寫死。另外在邊緣設(shè)備上，OpenVINO 或者 CoreML 的性能通常吊打普通 CPU 推理；如果是 Windows 平臺帶集顯DirectML 也是個容易被忽視的加速選項。

2.、像做手術(shù)一樣控制線程數(shù)（不要超配）

線程配置有兩個核心參數(shù)：intra-op（算子內(nèi)并行）和inter-op（算子間并行）。這兩個參數(shù)的設(shè)置必須參考機(jī)器的物理核心數(shù)以及你的負(fù)載特性。

import os, multiprocessing as mp, onnxruntime as ort
cores = mp.cpu_count() // 2 or 1 # conservative default
so = ort.SessionOptions()
so.intra_op_num_threads = cores
so.inter_op_num_threads = 1 # start low for consistent latency
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
sess = ort.InferenceSession("model.onnx", sess_options=so, providers=["CPUExecutionProvider"])

默認(rèn)的線程策略經(jīng)常會跟 NumPy、BLAS 庫甚至你的 Web Server 搶占資源，導(dǎo)致嚴(yán)重的線程爭用和長尾延遲。建議把 inter_op 設(shè)為 1（通常能獲得更穩(wěn)定的延遲），然后遍歷測試 intra_op（從 1 到物理核數(shù)），盯著p50和p95指標(biāo)找最佳平衡點(diǎn)，不要光看平均速度。

3、使用 IO Binding 規(guī)避內(nèi)存拷貝（GPU 必選項）

如果在 GPU 上跑推理，卻每次 run() 都把張量從 Device 拷回 Host再拷回 Device，利用 IO Binding 將輸入/輸出直接綁定在顯存上，復(fù)用這塊內(nèi)存。

import onnxruntime as ort
import numpy as np
sess = ort.InferenceSession("model.onnx", providers=["CUDAExecutionProvider"])
io = sess.io_binding()
# Example: preallocate on device via OrtValue (CUDA)
import onnxruntime as ort
x = np.random.rand(1, 3, 224, 224).astype(np.float32)
x_ort = ort.OrtValue.ortvalue_from_numpy(x, device_type="cuda", device_id=0)
io.bind_input(name=sess.get_inputs()[0].name, device_type="cuda", device_id=0, element_type=np.float32, shape=x.shape, buffer_ptr=x_ort.data_ptr())
io.bind_output(name=sess.get_outputs()[0].name, device_type="cuda", device_id=0)
sess.run_with_iobinding(io)
y_ort = io.get_outputs()[0] # still on device

這對于高頻請求特別重要，哪怕單次拷貝只耗費(fèi)幾毫秒，累積起來也是巨大的開銷，所以讓熱數(shù)據(jù)留在它該在的地方。

4、鎖定 Shape 或采用分桶策略

動態(tài) Shape 看起來很靈活，但它會阻礙 ORT 進(jìn)行激進(jìn)的算子融合和 Kernel 優(yōu)選。在導(dǎo)出 ONNX 時能固定 Shape 就盡量固定。

如果業(yè)務(wù)場景確實(shí)需要變長輸入，可以采用分桶（Bucketing）策略：

# pseudo: choose session by input shape
def get_session_for_shape(h, w):
if h <= 256 and w <= 256: return sess_256
if h <= 384 and w <= 384: return sess_384
return sess_fallback

比如在視覺任務(wù)中，把輸入限定在 224、256、384 這幾檔，創(chuàng)建對應(yīng)的 Session。哪怕只分兩三個桶，性能表現(xiàn)也比完全動態(tài) Shape 強(qiáng)得多。

5、開啟全圖優(yōu)化并驗證

這一步很簡單但容易被忽略。開啟 ORT_ENABLE_ALL，讓 ORT 幫你做算子融合、常量折疊和內(nèi)存規(guī)劃。

import onnxruntime as ort
so = ort.SessionOptions()
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
# optional: serialize the optimized model for inspection
so.optimized_model_filepath = "model.optimized.onnx"
sess = ort.InferenceSession("model.onnx", sess_options=so, providers=["CPUExecutionProvider"])

更少的算子意味著更少的 Kernel Launch 開銷和內(nèi)存帶寬壓力。建議導(dǎo)出一個 optimized_model_filepath，用 Netron 打開看看，確認(rèn) Conv+BN+ReLU 這種經(jīng)典組合是不是真的被融合成一個節(jié)點(diǎn)了，如果沒融那就是優(yōu)化鏈路上有問題。

6、CPU 推理？直接上量化

如果只能用 CPU，INT8 量化或者動態(tài)量化是提速神器。配合 CPU 的向量指令集能極大減少矩陣乘法的開銷。

from onnxruntime.quantization import quantize_dynamic, QuantType
quantize_dynamic(
model_input="model.onnx",
model_output="model.int8.onnx",
weight_type=QuantType.QInt8, # try QInt8 or QUInt8
extra_options={"MatMulConstBOnly": True}
)

然后加載量化后的模型：

import onnxruntime as ort
sess = ort.InferenceSession("model.int8.onnx", providers=["CPUExecutionProvider"])

對于 Transformer 類模型，動態(tài)量化通常能帶來 1.5 到 3 倍的加速且精度損失很小。不過需要先在真實(shí)數(shù)據(jù)上驗證，如果精度掉得厲害嘗試 Per-channel 量化或者只量化計算最密集的算子。

7、預(yù)熱、復(fù)用與 Micro-Batching

InferenceSession 的初始化開銷很大，屬于重資源對象。務(wù)必全局只創(chuàng)建一次，并且需要啟動后先跑幾次 Dummy Data 做預(yù)熱，把 Kernel Cache 和內(nèi)存池填好。

# app startup
sess = ort.InferenceSession("model.onnx", providers=["CUDAExecutionProvider"])
dummy = {sess.get_inputs()[0].name: np.zeros((1, 3, 224, 224), np.float32)}
for _ in range(3):
sess.run(None, dummy) # warms kernels, caches, memory arenas

如果是高并發(fā)場景不要一個個請求單獨(dú)跑，攢一個 Micro-batch（比如 2 到 8 個樣本）一起送進(jìn)去，能顯著提高 GPU 利用率（Occupancy）。

def infer_batch(batch):
inputs = np.stack(batch, axis=0).astype(np.float32, copy=False)
return sess.run(None, {sess.get_inputs()[0].name: inputs})[0]

調(diào)整 Batch Size 的時候，盯著p95 延遲和吞吐量看，找到那個甜點(diǎn)。

8、優(yōu)化前后處理：拒絕 Python 循環(huán)

很多時候大家抱怨模型慢，其實(shí)瓶頸在預(yù)處理和后處理。Python 的 for 循環(huán)處理像素或 logits 是絕對的性能殺手。所以保持?jǐn)?shù)組內(nèi)存連續(xù)，避免不必要的 astype 轉(zhuǎn)換盡量全部向量化。

import numpy as np
# Bad: repeated copies/conversions
# x = np.array(img).astype(np.float32) # realloc every time
# Better: reuse buffers and normalize in-place
buf = np.empty((1, 3, 224, 224), dtype=np.float32)
def preprocess(img, out=buf):
# assume img is already CHW float32 normalized upstream
np.copyto(out, img, casting="no") # no implicit cast
return out
# Post-process with NumPy ops, not Python loops
def topk(logits, k=5):
idx = np.argpartition(logits, -k, axis=1)[:, -k:]
vals = np.take_along_axis(logits, idx, axis=1)
order = np.argsort(-vals, axis=1)
return np.take_along_axis(idx, order, axis=1), np.take_along_axis(vals, order, axis=1)

幾個多余的 .astype() 就能吃掉好幾毫秒，這點(diǎn)在低延遲場景下非常致命。

基準(zhǔn)測試模板

這是一個簡單的 Benchmarking 腳本，改改就能用，別靠感覺優(yōu)化要用數(shù)據(jù)來進(jìn)行對比：

import time, statistics as stats
import numpy as np, onnxruntime as ort
def bench(sess, x, iters=100, warmup=5):
name = sess.get_inputs()[0].name
for _ in range(warmup):
sess.run(None, {name: x})
times = []
for _ in range(iters):
t0 = time.perf_counter()
sess.run(None, {name: x})
times.append((time.perf_counter() - t0) * 1e3)
return {
"p50_ms": stats.median(times),
"p95_ms": sorted(times)[int(0.95 * len(times)) - 1],
"min_ms": min(times),
"max_ms": max(times)
}
# Example usage
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
so = ort.SessionOptions(); so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
sess = ort.InferenceSession("model.onnx", sess_options=so, providers=providers)
x = np.random.rand(1, 3, 224, 224).astype(np.float32)
print(bench(sess, x))

總結(jié)

做低延遲推理沒有什么黑科技，全是細(xì)節(jié)。選對 Provider，別亂開線程，減少內(nèi)存拷貝，固定 Shape，激進(jìn)地做圖融合，最后把 Python 代碼洗干凈。哪怕只落實(shí)其中兩三點(diǎn)，性能提升也是肉眼可見的。

https://avoid.overfit.cn/post/aa489c6b429641b9b1a1a3e4a3e4ce1d

作者：Modexa