TensorRT Engine 文件

.engine 和 .plan 文件格式是由 NVIDIA 的 TensorRT 用於高性能深度學習推理的格式。

功能

TensorRT 引擎文件包含為特定 GPU 硬件和驅動版本優化的序列化執行計劃。這些優化帶來了：

• 低延遲：比標準框架更快的推理速度。
• 高吞吐量：在相同硬件上每秒完成更多的推理任務。

優化技術

技術	描述
層融合 (Layer Fusion)	將多個操作合併為一個內核（例如：Conv + ReLU）。
精度校準 (Precision Calibration)	支持 FP32、FP16 和 INT8 量化以提高效率。
內核自動調優 (Kernel Auto-tuning)	為目標硬件選擇最佳的內核實現。
內存優化 (Memory Optimization)	通過管理中間張量內存來減少內存占用。

創建 Engine 文件

從 ONNX (推薦)

import tensorrt as trt

# 創建構建器
logger = trt.Logger(trt.Logger.INFO)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

# 解析 ONNX 模型
with open("model.onnx", "rb") as model:
    parser.parse(model.read())

# 構建並序列化引擎
config = builder.create_builder_config()
engine = builder.build_engine(network, config)

with open("model.engine", "wb") as f:
    f.write(engine.serialize())

從 PyTorch (Torch-TensorRT)

import torch
import torch_tensorrt

# 編譯為 TensorRT 引擎
trt_model = torch_tensorrt.compile(
    model,
    inputs=[torch_tensorrt.Input((1, 3, 224, 224), dtype=torch.half)],
    enabled_precisions={torch.half},  # FP16
    workspace_size=1 << 30
)

# 保存引擎
torch.jit.save(trt_model, "model_torchtrt.engine")

硬件綁定

TensorRT 引擎文件與其構建時的特定 GPU 架構綁定。在 Tesla V100 上構建的引擎無法在 GeForce RTX 3090 上運行。

加載和使用 Engine

import tensorrt as trt

# 加載引擎
with open("model.engine", "rb") as f, trt.Runtime(trt.Logger(trt.Logger.INFO)) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())

# 創建推理上下文
context = engine.create_execution_context()

小貼士

構建優化的引擎可能需要幾分鐘。對於生產環境，建議預構建引擎並進行快取，以減少應用啟動時間。