TensorRT 使用方式与典型场景

TensorRT 使用方式与典型场景

流程固定为四步：导出 ONNX、解析并构建 engine、序列化保存、反序列化执行。第一条示例用 Python API 从 ONNX 直接构建 FP16 engine，并声明动态 shape 范围。

import tensorrt as trt

logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

with open('model.onnx', 'rb') as f:
    if not parser.parse(f.read()):
        for i in range(parser.num_errors):
            print(parser.get_error(i))
        raise SystemExit('ONNX 解析失败')

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
config.set_flag(trt.BuilderFlag.FP16)

# 动态 shape：声明输入范围，opt 档务必贴近线上真实 batch
profile = builder.create_optimization_profile()
profile.set_shape('input',
                  (1, 3, 224, 224),
                  (8, 3, 224, 224),
                  (32, 3, 224, 224))
config.add_optimization_profile(profile)

plan = builder.build_serialized_network(network, config)
with open('model.engine', 'wb') as f:
    f.write(plan)

第二条示例是加载 engine 并执行推理：set_input_shape 必须落在 profile 范围内，张量地址要与 dtype、显存对齐一致。

import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

logger = trt.Logger(trt.Logger.INFO)
runtime = trt.Runtime(logger)
with open('model.engine', 'rb') as f:
    engine = runtime.deserialize_cuda_engine(f.read())

context = engine.create_execution_context()
context.set_input_shape('input', (1, 3, 224, 224))

shape = context.get_tensor_shape('input')
d_input = cuda.mem_alloc(int(np.prod(shape)) * 4)
d_output = cuda.mem_alloc(1000 * 4)
context.set_tensor_address('input', int(d_input))
context.set_tensor_address('output', int(d_output))

stream = cuda.Stream()
host_input = np.random.rand(*shape).astype(np.float32)
cuda.memcpy_htod_async(d_input, host_input, stream)
context.execute_async_v3(stream_handle=stream.handle)
cuda.memcpy_dtoh_async(host_output, d_output, stream)
stream.synchronize()

不想写代码的话，命令行工具 trtexec 已经覆盖了绝大多数构建与压测需求，参数与 Python API 基本一一对应。

trtexec --onnx=model.onnx \
  --saveEngine=model.engine \
  --fp16 \
  --minShapes=input:1x3x224x224 \
  --optShapes=input:8x3x224x224 \
  --maxShapes=input:32x3x224x224 \
  --workspace=2048 \
  --iterations=200 --avgRuns=50

要开 INT8，就得提供一个校准器，用几百张贴近线上分布的无标注样本统计激活范围，并把校准缓存落盘复用。

class Calibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, batches, batch_size=8):
        super().__init__()
        self.batches = iter(batches)
        self.batch_size = batch_size

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        batch = next(self.batches, None)
        if batch is None:
            return None
        cuda.memcpy_htod(self.d_input, batch)
        return [int(self.d_input)]

    def read_calibration_cache(self):
        return None

    def write_calibration_cache(self, cache):
        with open('calib.cache', 'wb') as f:
            f.write(cache)

适合的场景：图像分类、检测、分割、OCR、人脸、语音，以及推荐模型的排序层，尤其是 batch 相对固定、延迟指标敏感的在线服务。不适合的场景：模型结构还在频繁改动（每改一次都要重编译）、图里有大量动态控制流且难以用 plugin 表达、需要跨多家硬件厂商（这时 ONNX Runtime 或原生框架更合适），以及训练本身。