Stories

Detail Return Return

TensorRT 和 ONNX Runtime 推理優化實戰:10 個降低延遲的工程技巧 - Stories Detail

模型速度的瓶頸往往不在算法本身。幾毫秒的優化累積起來就能讓用户感受到明顯的性能提升。下面這些技術都是在生產環境跑出來的經驗,不需要重構代碼實施起來也相對簡單並且效果顯著。

固定輸入形狀,越早告訴運行時越好

動態形狀用起來方便但對性能不友好。TensorRT 和 ONNX Runtime 在處理固定形狀時能做更激進的優化。

TensorRT 這邊,構建引擎時最好圍繞實際使用的 min/opt/max 設置 optimization profile,生產環境儘量讓所有請求都落在 opt 範圍。ONNX Runtime 可以直接導出固定維度的模型,比如 1×3×224×224。確實需要動態性的話,也要把範圍控制得足夠緊。

 # TensorRT: build with a tight optimization profile  
 profile = builder.create_optimization_profile()  
 profile.set_shape("input", (1,3,224,224), (1,3,224,224), (1,3,224,224))  
 config.add_optimization_profile(profile)

這樣kernel 選擇、內存規劃、算子融合在形狀確定的情況下都能做得更徹底。

啓動前把該熱的都熱一遍

冷啓動的開銷藏在各個角落:驅動初始化、page fault、lazy allocation。服務啓動和重啓後跑幾輪 warmup,把這些一次性成本提前消化掉。

 # ONNX Runtime warmup + pinned buffers  
import onnxruntime as ort, numpy as np  

so = ort.SessionOptions()  
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL  
sess = ort.InferenceSession("model.onnx", sess_options=so, providers=["CUDAExecutionProvider"])  

x = np.random.randn(1,3,224,224).astype(np.float32)  
for _ in range(8):  # small loop to populate caches/contexts  
     sess.run(None, {"input": x})

warmup 的形狀一定要和線上一模一樣。如果服務多種 batch size,每個都得過一遍。

I/O binding 配合 pinned memory 減少拷貝

Host 和 device 之間來回搬數據是 tail latency 的大敵。buffer 綁定一次,後面反覆用就行了。

 # ONNX Runtime I/O binding example  
io = sess.io_binding()  
x = np.random.randn(1,3,224,224).astype(np.float32)  

# Upload once & bind  
io.bind_cpu_input("input", x)       # or bind to CUDA device via OrtValue  
io.bind_output("logits", device_type="cuda")  

sess.run_with_iobinding(io)  
 out = io.copy_outputs_to_cpu()[0]   # pull back only when you must

GPU 流量大的場景,host 端內存用 page-locked(pinned)能讓 H2D/D2H 傳輸快不少。本質上是把多次小開銷合併成一次前置成本,allocator 也不用頻繁介入。

精度降低不一定掉點

現在的 GPU 對 FP16 支持很好,服務器 CPU 和 NPU 上 INT8 的收益也越來越明顯。只要精度守得住,延遲的改善非常直接。

TensorRT 開 FP16 就是一個 flag 的設置:

config.set_flag(trt.BuilderFlag.FP16)

。但是INT8 需要校準,要拿代表性數據跑一遍生成 per-channel scale。ONNX Runtime 可以用 TensorRT EP 或者直接加載量化後的模型。

 # TensorRT FP16 (and INT8 if you have a calibrator)  
 config.set_flag(trt.BuilderFlag.FP16)  
 # config.set_flag(trt.BuilderFlag.INT8)  
 # config.int8_calibrator = calibrator

這裏可以先量化最慢的幾個子圖,比如 embedding 層或者 attention block,不用一上來就全模型量化。

圖優化可以開到最高檔,但要驗證數值

讓運行時自己去融合算子、選更優的 kernel,這個收益基本是白來的。

 # ONNX Runtime optimizations + EPs  
so = ort.SessionOptions()  
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL  
providers = [  
  ("TensorrtExecutionProvider", {"trt_fp16_enable": True}),  
  "CUDAExecutionProvider",  
  "CPUExecutionProvider"  
]  
 sess = ort.InferenceSession("model.onnx", sess_options=so, providers=providers)

但是需要注意的是融合會改變計算順序,數值可能有細微漂移。開啓前後要跑個 tolerance check,確保輸出沒問題。

micro-batch 在 GPU 上效果明顯

單條請求跑推理簡單,但硬件利用率往往上不去。打包成 4-8 個請求一起跑,能在保持低延遲的同時提升吞吐。

關鍵是 batching window 要夠小,比如 2-5ms,不然 p95 會飆。micro-batch 的大小最好和前面 optimization profile 設置的尺寸對齊。

不過如果 SLA 本身就很緊(p50 要求 5ms 以內),micro-batch 帶來的收益可能不如下面要説的 CUDA Graph。

CUDA Graph 消除 kernel launch 開銷

小模型或者調用頻繁的 graph,kernel launch 的開銷會很明顯。CUDA Graph 能把整個推理過程錄製下來,replay 時幾乎沒有 CPU 開銷。

TensorRT 在形狀固定的情況下可以直接啓用,只需要warmup 一次,後面就一直跑 captured graph。

這裏可以理解成在 GPU driver 層面把推理編譯成一個可重放的宏。

ONNX Runtime 線程設置有講究

ONNX Runtime 暴露了幾個線程相關的參數,對 CPU 和混合負載的 tail latency 影響挺大。

 so = ort.SessionOptions()  
 so.intra_op_num_threads = 1   # one thread per operator often stabilizes latency  
 so.inter_op_num_threads = 1   # avoid oversubscription; raise carefully if parallel graphs

Execution Provider 的選擇也很重要:

GPU 場景優先級是 TensorRT EP → CUDA EP → CPU EP fallback。純 CPU 跑的話 OpenMP 或者 DNNL/MKL build 配合合理的線程池設置效果最好。邊緣設備上 Intel 的盒子可以試試 OpenVINO EP。

把預處理後處理從 GIL 裏挪出去

Python 的膠水代碼經常成為隱藏的性能殺手。能用 NumPy 向量化就別寫循環,能用 Numba 或者推到 CUDA/CuPy 上更好。熱路徑裏的 transform 最好提前編譯好。如果要併發處理請求,worker pool 的規模要和運行時的線程數配合好。

 # Example: pre-allocate and reuse buffers to dodge Python overheads  
import numpy as np  

class Preprocessor:  
    def __init__(self, shape=(1,3,224,224)):  
        self.buf = np.empty(shape, dtype=np.float32)  

    def __call__(self, img):  
        # write into self.buf in-place; no fresh allocations  
        np.copyto(self.buf, img)  
        self.buf /= 255.0  
         return self.buf

這裏的判斷標準很簡單,每個請求都會跑的代碼,問問能不能預分配、向量化或者緩存起來。

Session、Engine、Buffer 都只建一次

每個請求新建一個

trt.ICudaEngine

onnxruntime.InferenceSession

基本等於自殺。output array 每次重新分配也一樣。

正確做法是進程啓動時就加載好 singleton session/engine,每個 worker 維護一兩個 CUDA stream,buffer pool 按 shape 和 dtype 索引。

 # Simple singleton-ish loader  
class Model:  
    _sess = None  
    _io = None  
    @classmethod  
    def get(cls):  
        if cls._sess is None:  
            so = ort.SessionOptions()  
            so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL  
            cls._sess = ort.InferenceSession("model.onnx", sess_options=so,  
                          providers=["CUDAExecutionProvider"])  
            cls._io = cls._sess.io_binding()  
         return cls._sess, cls._io

這樣做的好處是穩定,p95 不會因為 allocator 和 initializer 出現在熱路徑而突然炸掉。

一個完整的 GPU 推理骨架

下面的代碼是把前面幾個關鍵技術串起來:

 import onnxruntime as ort, numpy as np  

def make_session(path):  
    so = ort.SessionOptions()  
    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL  
    providers = [("TensorrtExecutionProvider", {"trt_fp16_enable": True}),  
                 "CUDAExecutionProvider", "CPUExecutionProvider"]  
    return ort.InferenceSession(path, sess_options=so, providers=providers)  

class Runner:  
    def __init__(self, model_path, shape=(1,3,224,224)):  
        self.sess = make_session(model_path)  
        self.shape = shape  
        self.io = self.sess.io_binding()  
        self._warmup()  

    def _warmup(self, iters=8):  
        x = np.random.randn(*self.shape).astype(np.float32)  
        self.io.bind_cpu_input("input", x)  
        self.io.bind_output("logits", device_type="cuda")  
        for _ in range(iters):  
            self.sess.run_with_iobinding(self.io)  
        self.io.clear_binding_inputs()  # ready for real runs  

    def run(self, x_np: np.ndarray):  
        # assumes x_np matches self.shape; in production, validate or clamp  
        self.io.bind_cpu_input("input", x_np)  
        self.io.bind_output("logits", device_type="cuda")  
        self.sess.run_with_iobinding(self.io)  
        return self.io.copy_outputs_to_cpu()[0]  

# usage  
runner = Runner("model.onnx")  
batch = np.random.randn(1,3,224,224).astype(np.float32)  
 probs = runner.run(batch)

這個代碼已經包含了圖優化、I/O binding 和 warmup。後面再加上 CUDA Graph、micro-batch 和固定 shape,能把延遲壓到很低,基本上拿來就可以用了

幾個容易踩的坑

延遲指標一定要看 p50/p90/p95,別隻盯平均值。真正的問題都藏在 tail 裏。API 層面最好把 shape 和 dtype 固定下來,或者至少讓調用方知道優化過的範圍。這樣生產請求才能穩定落在最優路徑上。

開了融合或量化之後,精度的自動化迴歸測試必不可少。

總結

低延遲不靠黑科技就是一堆小優化疊起來:形狀固定、減少拷貝、更好的 kernel、graph capture、運行時零意外。每個單拎出來可能只省幾毫秒,但加起來用户就能感受到"快"。

https://avoid.overfit.cn/post/494ca93b9c184407936ef7b6bd16e15e

作者:Syntal

Add a new Comments

Some HTML is okay.