在深度學習落地過程中,有一個常見的誤區:一旦推理速度不達標,大家的第一反應往往是拿着模型開到,比如:做剪枝、搞蒸餾、甚至犧牲精度換小模型。
實際上生產環境中的 Python 推理鏈路隱藏着巨大的“工程紅利”。很多時候你的模型本身並不慢,慢的是低效的數據搬運、混亂的線程爭用以及不合理的 Runtime 默認配置。在不改變模型精度的情況下,僅靠ONNX Runtime (ORT) 的工程特性,往往就能從現有技術棧中“摳”出驚人的性能提升。
以下是 8 個經過實戰驗證的低延遲優化策略,專治各種“莫名其妙的慢”。
1、 明確指定 Execution Provider 及其順序
ORT 會嚴格按照你傳入的
providers
列表順序進行嘗試。把最快的放在第一位,並且儘量避免它靜默回退(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 是有成本的,如果環境裏有 TensorRT 就優先用,沒有就降級到 CUDA,最後才是 CPU。把這個路徑寫死。另外在邊緣設備上,OpenVINO 或者 CoreML 的性能通常吊打普通 CPU 推理;如果是 Windows 平台帶集顯DirectML 也是個容易被忽視的加速選項。
2.、像做手術一樣控制線程數(不要超配)
線程配置有兩個核心參數:intra-op(算子內並行)和 inter-op(算子間並行)。這兩個參數的設置必須參考機器的物理核心數以及你的負載特性。
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"])
默認的線程策略經常會跟 NumPy、BLAS 庫甚至你的 Web Server 搶佔資源,導致嚴重的線程爭用和長尾延遲。建議把
inter_op
設為 1(通常能獲得更穩定的延遲),然後遍歷測試
intra_op
(從 1 到物理核數),盯着 p50 和 p95 指標找最佳平衡點,不要光看平均速度。
3、使用 IO Binding 規避內存拷貝(GPU 必選項)
如果在 GPU 上跑推理,卻每次
run()
都把張量從 Device 拷回 Host再拷回 Device,利用 IO Binding 將輸入/輸出直接綁定在顯存上,複用這塊內存。
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
這對於高頻請求特別重要,哪怕單次拷貝只耗費幾毫秒,累積起來也是巨大的開銷,所以讓熱數據留在它該在的地方。
4、鎖定 Shape 或採用分桶策略
動態 Shape 看起來很靈活,但它會阻礙 ORT 進行激進的算子融合和 Kernel 優選。在導出 ONNX 時能固定 Shape 就儘量固定。
如果業務場景確實需要變長輸入,可以採用分桶(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
比如在視覺任務中,把輸入限定在 224、256、384 這幾檔,創建對應的 Session。哪怕只分兩三個桶,性能表現也比完全動態 Shape 強得多。
5、開啓全圖優化並驗證
這一步很簡單但容易被忽略。開啓
ORT_ENABLE_ALL
,讓 ORT 幫你做算子融合、常量摺疊和內存規劃。
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 開銷和內存帶寬壓力。建議導出一個
optimized_model_filepath
,用 Netron 打開看看,確認 Conv+BN+ReLU 這種經典組合是不是真的被融合成一個節點了,如果沒融那就是優化鏈路上有問題。
6、CPU 推理?直接上量化
如果只能用 CPU,INT8 量化或者動態量化是提速神器。配合 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 類模型,動態量化通常能帶來 1.5 到 3 倍的加速且精度損失很小。不過需要先在真實數據上驗證,如果精度掉得厲害嘗試 Per-channel 量化或者只量化計算最密集的算子。
7、預熱、複用與 Micro-Batching
InferenceSession
的初始化開銷很大,屬於重資源對象。務必全局只創建一次,並且需要啓動後先跑幾次 Dummy Data 做預熱,把 Kernel Cache 和內存池填好。
# app startup
sess=ort.InferenceSession("model.onnx", providers=["CUDAExecutionProvider"])
dummy= {sess.get_inputs()[0].name: np.zeros((1, 3, 224, 224), np.float32)}
for_inrange(3):
sess.run(None, dummy) # warms kernels, caches, memory arenas
如果是高併發場景不要一個個請求單獨跑,攢一個 Micro-batch(比如 2 到 8 個樣本)一起送進去,能顯著提高 GPU 利用率(Occupancy)。
definfer_batch(batch):
inputs=np.stack(batch, axis=0).astype(np.float32, copy=False)
returnsess.run(None, {sess.get_inputs()[0].name: inputs})[0]
調整 Batch Size 的時候,盯着 p95 延遲和吞吐量看,找到那個甜點。
8、優化前後處理:拒絕 Python 循環
很多時候大家抱怨模型慢,其實瓶頸在預處理和後處理。Python 的
for
循環處理像素或 logits 是絕對的性能殺手。所以保持數組內存連續,避免不必要的
astype
轉換儘量全部向量化。
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()
就能吃掉好幾毫秒,這點在低延遲場景下非常致命。
基準測試模板
這是一個簡單的 Benchmarking 腳本,改改就能用,別靠感覺優化要用數據來進行對比:
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))
總結
做低延遲推理沒有什麼黑科技,全是細節。選對 Provider,別亂開線程,減少內存拷貝,固定 Shape,激進地做圖融合,最後把 Python 代碼洗乾淨。哪怕只落實其中兩三點,性能提升也是肉眼可見的。
https://avoid.overfit.cn/post/aa489c6b429641b9b1a1a3e4a3e4ce1d
作者:Modexa