博客 / 詳情

返回

ONNX Runtime Python 推理性能優化:8 個低延遲工程實踐

在深度學習落地過程中,有一個常見的誤區:一旦推理速度不達標,大家的第一反應往往是拿着模型開到,比如:做剪枝、搞蒸餾、甚至犧牲精度換小模型。

實際上生產環境中的 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 到物理核數),盯着 p50p95 指標找最佳平衡點,不要光看平均速度。

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

user avatar
0 位用戶收藏了這個故事!

發佈 評論

Some HTML is okay.