模型速度的瓶頸往往不在算法本身。幾毫秒的優化累積起來就能讓用户感受到明顯的性能提升。下面這些技術都是在生產環境跑出來的經驗,不需要重構代碼實施起來也相對簡單並且效果顯著。
固定輸入形狀,越早告訴運行時越好
動態形狀用起來方便但對性能不友好。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