如何使用execute_async_v3()配置TensorRT推理上下文？

execute_async_v3() 替代弃用的execute_async_v2()后，核心变化是不再通过缓冲区列表传递内存地址，而是需要显式为每个输入/输出张量绑定设备内存地址。以下是完整实现流程及示例：

1. 基础准备 #

假设你已完成模型转换得到engine对象，并创建了执行上下文context（通过ICudaEngine.create_execution_context()）。

2. 动态输入形状配置（动态形状引擎必填） #

如果你的引擎支持动态batch/尺寸，必须先设置输入张量的实际形状：

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

# 示例输入形状：(batch_size, channel, height, width)
input_shape = (1, 3, 224, 224)
# 替换为你的输入张量名称，可通过engine.get_tensor_name(0)获取
context.set_input_shape("input", input_shape)

3. 分配输入输出设备内存 #

计算张量所需内存大小，分配CUDA设备内存：

# 获取输入输出张量名称
input_name = context.engine.get_tensor_name(0)
output_name = context.engine.get_tensor_name(1)

# 计算字节数：总元素数 × 数据类型字节大小
input_byte_size = trt.volume(input_shape) * trt.nptype(trt.float32).itemsize
output_shape = context.get_tensor_shape(output_name)
output_byte_size = trt.volume(output_shape) * trt.nptype(trt.float32).itemsize

# 分配设备内存
input_device_ptr = cuda.mem_alloc(input_byte_size)
output_device_ptr = cuda.mem_alloc(output_byte_size)

# 准备主机端输入数据（示例数据）
host_input = np.random.randn(*input_shape).astype(np.float32)

4. 绑定张量地址到执行上下文 #

这是execute_async_v3()的核心步骤，必须为每个输入输出张量绑定对应的设备内存地址：

context.set_tensor_address(input_name, input_device_ptr)
context.set_tensor_address(output_name, output_device_ptr)

5. 异步推理执行 #

通过CUDA流完成异步数据拷贝与推理：

# 创建CUDA流
stream = cuda.Stream()

# 主机→设备 异步拷贝输入数据
cuda.memcpy_htod_async(input_device_ptr, host_input, stream)

# 执行异步推理，仅需传入CUDA流句柄
success = context.execute_async_v3(stream.handle)
if not success:
    raise RuntimeError("推理执行失败")

# 设备→主机 异步拷贝输出结果
host_output = np.empty(output_shape, dtype=np.float32)
cuda.memcpy_dtoh_async(host_output, output_device_ptr, stream)

# 同步流，等待所有异步操作完成
stream.synchronize()

# 处理输出结果
print("推理输出形状:", host_output.shape)

关键注意事项 #

• 张量名称匹配：必须确保set_tensor_address()使用的名称与引擎中输入输出张量的名称完全一致，可通过以下代码遍历确认：

  for i in range(context.engine.num_io_tensors):
      name = context.engine.get_tensor_name(i)
      mode = context.engine.get_tensor_mode(name)
      print(f"{'输入' if mode == trt.TensorIOMode.INPUT else '输出'}张量: {name}")
  

• 静态形状引擎：若引擎为静态形状，无需调用set_input_shape()，但仍需绑定张量地址。
• 内存校验：确保设备内存分配尺寸与张量形状匹配，避免因内存不足或尺寸不匹配导致的推理失败。
• 流同步：异步操作后必须同步CUDA流，否则主机端获取的输出数据可能不完整。
• 错误检查：execute_async_v3()返回布尔值，务必检查返回结果是否为True，避免静默失败。

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