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()

小贴士

构建优化的引擎可能需要几分钟。对于生产环境，建议预构建引擎并进行缓存，以减少应用启动时间。