Windows下使用trtexec运行TensorRT推理如何导出二进制输出结果

官方发行版的trtexec工具原生不支持直接导出二进制格式的推理结果，自带的--exportOutput参数仅支持输出JSON格式的单条推理结果，确实不适配Image2Image这类需要批量处理大尺寸输出张量的场景。

优先推荐无需编译的Python方案，仅需少量代码即可完全对齐trtexec的FP16推理逻辑，且直接支持二进制输入/输出：

• 前置准备：
  1. 执行原有trtexec命令时新增--saveEngine=model_fp16.engine参数，导出构建完成的FP16序列化引擎文件，避免后续重复构建
  2. 用pip安装对应版本的依赖包：pip install tensorrt==<你的TensorRT版本号> pycuda onnxruntime-gpu
• TensorRT Python API推理代码示例：

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

# 加载序列化FP16引擎
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
with open("model_fp16.engine", "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()

# 替换为你自己的输入输出名称、shape、dtype
input_shape = (1, 3, 512, 512)
input_dtype = np.float16
output_shape = (1, 3, 512, 512)
output_dtype = np.float16

# 分配设备内存
d_input = cuda.mem_alloc(int(np.prod(input_shape) * np.dtype(input_dtype).itemsize))
d_output = cuda.mem_alloc(int(np.prod(output_shape) * np.dtype(output_dtype).itemsize))
bindings = [int(d_input), int(d_output)]
stream = cuda.Stream()

# 批量处理所有二进制输入blob
input_dir = "./input_blobs"
output_dir = "./output_blobs"
os.makedirs(output_dir, exist_ok=True)

for blob_file in os.listdir(input_dir):
    if not blob_file.endswith(".blob"):
        continue
    # 读取二进制输入
    input_data = np.fromfile(os.path.join(input_dir, blob_file), dtype=input_dtype).reshape(input_shape)
    # 拷贝输入到GPU、执行推理、拷贝结果回CPU
    cuda.memcpy_htod_async(d_input, input_data, stream)
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    output_data = np.empty(output_shape, dtype=output_dtype)
    cuda.memcpy_dtoh_async(output_data, d_output, stream)
    stream.synchronize()
    # 输出保存为二进制blob
    output_data.tofile(os.path.join(output_dir, f"res_{blob_file}"))

• 如果你嫌TensorRT Python API的内存分配逻辑繁琐，也可以替换为ONNX Runtime + TensorRT执行提供器的方案，开启FP16精度后推理结果和trtexec完全一致，代码更简洁：

import onnxruntime as ort
import numpy as np
import os

# 初始化ONNX Runtime会话，开启TensorRT FP16加速
sess = ort.InferenceSession(
    "your_model.onnx",
    providers=["TensorrtExecutionProvider"],
    provider_options=[{"trt_fp16_enable": True, "trt_engine_cache_enable": True}]
)
input_dir = "./input_blobs"
output_dir = "./output_blobs"
os.makedirs(output_dir, exist_ok=True)

for blob_file in os.listdir(input_dir):
    if not blob_file.endswith(".blob"):
        continue
    input_data = np.fromfile(os.path.join(input_dir, blob_file), dtype=np.float16).reshape(1,3,512,512)
    output_data = sess.run(None, {"your_input_name": input_data})[0]
    output_data.tofile(os.path.join(output_dir, f"res_{blob_file}"))

如果必须使用trtexec工具链，可以修改TensorRT官方开源的trtexec样例代码，新增自定义参数控制输出张量直接写入二进制文件即可，不过需要在Windows下配置CUDA、TensorRT、CMake的编译环境，实现成本更高，仅适合有定制化工具链需求的场景。

内容的提问来源于stack exchange，提问作者ma0ho