Tensorrt python api 10.11.0笔记
关于Tensorrt的python api文档阅读翻译加总结
文档源地址
Overview
Getting started with TensorRT
Installation(安装)
安装可参考:官方地址
Samples
关于样例的内容可参考:样例地址
Operator Documentation
有关更多信息（包括示例），请参阅 TensorRT Operator’s Reference documentation
Installing cuda-python
虽然 TensorRT Python API 不需要，但 cuda-python 用于多个示例。有关安装说明，请参阅 CUDA Python 安装文档 。
Core Concepts
TensorRT Workflow工作流
Class Overview类概述
Logger(记录器)
Tensorrt提供一个tensorrt.Logger实现,基本功能都有,但是可以重写获取更高级的功能.
Parsers(解析器)
Parasers用于填充tensorrt.INetworkDefinition(来自深度学习训练的模型)
Newwork(网络)
tensorrt.INetworkDefinition表示一个计算图，为了填充网络，TensorRT提供了一套适用于各种深度学习框架的解析器，您也可以使用网络API手动填充网络
Builder(构建器)
tensorrt.Builder用于构建tensorrt.ICudaEngine
为此必须提供一个tensorrt.INetworkDefinition
Engine and Context(引擎和上下文)
tensorrt.ICudaEngine是TensorRT的输出，生成可执行推理的tensorrt.IExecutionContext.
Writing custom operators with TensorRT python plugins(使用TensorRT Python插件编写自定义运算符)
本指南展示了如何实现和包装定义插件行为的 Python 函数，以便将其作为自定义算子添加到网络中。
Composition of a plugin(plugin的组成)
需要定义两个函数，由tensorrt.plugin提供的装饰器包装
1、tensorrt.plugin.register()
tensorrt.plugin.register()返回输出张量的形状和类型特征。函数签名还定义了输入张量以及插件运行所需的任何属性。
2、tensorrt.plugin.impl()ortensorrt.plugin.aot_impl()
tensorrt.plugin.impl():JIT计算定义(A Just-in-Time compute definition)
tensorrt.plugin.aot_impl()AOT计算定义(A Ahead-of-Time compute definition)
Example: Circular padding plugin(圆形填充插件)
具体内容可见此官方链接
使用经典JIT实现（一般需要register和impl两个函数）
TensorRT 允许通过自定义插件扩展支持的算子（如这里的循环填充）。插件需要定义两部分关键信息：
形状描述：告知 TensorRT 输入输出的形状、数据类型等元信息（由 circ_pad_plugin_desc 函数实现）。
import tensorrt.plugin as trtp
import numpy.typing as npt@trtp.register("example::circ_pad_plugin")
def circ_pad_plugin_desc( # 注册插件名称()inp0: trtp.TensorDesc, # 输入张量的描述（形状、数据类型等）pads: npt.NDArray[np.int32]# 填充参数（各维度的前后填充量）
) -> trtp.TensorDesc: # 返回输出张量的描述ndim = inp0.ndimout_desc = inp0.like()for i in range(np.size(pads) // 2):out_desc.shape_expr[ndim - i - 1] += int(pads[i * 2] + pads[i * 2 + 1])return out_desc
cicular padding plugin的具体实现如下
import tensorrt.plugin as trtp@trtp.impl("example::circ_pad_plugin")# 注册插件实现（名称需与形状描述函数一致）
def circ_pad_plugin_impl(inp0: trtp.Tensor,# 输入张量（来自 TensorRT 的输入绑定）pads: npt.NDArray[np.int32],# 填充参数（各维度的前后填充量）outputs: Tuple[trtp.Tensor],# 输出张量元组（可能包含多个输出，此处仅一个）stream: int # CUDA 流句柄（用于计算同步）
) -> None:inp_t = torch.as_tensor(inp0, device="cuda")# 将输入张量转为 PyTorch CUDA 张量out_t = torch.as_tensor(outputs[0], device="cuda")# 将输出张量转为 PyTorch CUDA 张量out = torch.nn.functional.pad(inp_t, pads.tolist(), mode="circular")# 循环填充out_t.copy_(out)# 将填充结果复制到输出张量（共享内存时可省略，但显式复制更安全）
关于以上两端代码的内容，总结下来为：这两段代码分别对应 TensorRT 自定义插件开发中的两个核心部分：插件描述注册（Shape Inference）和插件计算实现（Execution），主要区别如下：
1.功能定位不同
第一段（@trtp.register装饰的函数）：
负责定义插件的形状推断逻辑（Shape Inference）。
作用是在 TensorRT 构建网络（如解析模型时）时，根据输入张量的描述（TensorDesc）和插件参数（pads），计算输出张量的形状（如维度大小），确保 TensorRT 能正确分配内存、优化网络结构。
第二段（@trtp.impl装饰的函数）：
负责定义插件的实际计算逻辑（Execution）。
作用是在模型推理时，根据输入张量的数据、插件参数（pads）和 CUDA 流（stream），完成具体的计算操作（如循环填充），并将结果写入输出张量。
2.执行截断不同
第一段（形状推断）：
在 TensorRT 网络构建阶段（如解析 ONNX 模型、优化引擎时）执行，仅需要元信息（形状、数据类型），不需要实际数据。
第二段（计算实现）：
在模型推理阶段（引擎运行时）执行，需要操作真实的输入 / 输出数据，并通过 CUDA 流管理计算同步。
Providing an Ahead-of-time(AOT)implement (提供预先AOT实现)
与tensorrt.plugin.impl()提供的JIT实现相比，AOT实现具有以下优势：
1、为插件构建具有 JIT 计算功能的 TRT 引擎需要其 tensorrt.plugin.register() 和 tensorrt.plugin.impl() Python 定义在运行时均存在。使用 AOT 实现时，插件将完全嵌入到引擎中，因此运行时无需用户提供插件定义。
2、插件层在运行时独立于 Python。这意味着，如果 TRT 引擎仅包含 AOT 插件，它可以在标准 TRT 运行时上执行，就像使用已编译的 C++ 插件的引擎一样，例如通过 trtexec 执行。
tensorrt.plugin.aot_impl()可用于插件的AOT实现。使用OpenAI Triton内核定义一个AOT实现：
import triton
import triton.language as tl@triton.jit
# Triton内核，定义循环填充的具体计算逻辑
def circ_pad_kernel(# input tensorX,#输入数据，存储在GPU内存中# extra scalar args in between input and output tensors 填充参数和维度信息（如输入/输出的各维度大小，总元素数）all_pads_0,all_pads_2,all_pads_4,all_pads_6,orig_dims_0,orig_dims_1,orig_dims_2,orig_dims_3,Y_shape_1,Y_shape_2,Y_shape_3,X_len,Y_len,# output tensorY, #填充后的结果（需写入 GPU 内存）BLOCK_SIZE: tl.constexpr,#Triton 内核的线程块大小
):pid = tl.program_id(0) #通过 tl.program_id(0) 获取当前线程块的 ID（pid）i = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) #计算每个线程的全局索引imask_y = i < Y_len
#将全局索引 i 分解为输出张量的多维坐标 (i0, i1, i2, i3)（对应输出的各维度位置）i3 = i % Y_shape_3i2 = (i // Y_shape_3) % Y_shape_2i1 = (i // Y_shape_3 // Y_shape_2) % Y_shape_1i0 = i // Y_shape_3 // Y_shape_2 // Y_shape_1
#将输出的位置 (i0, i1, i2, i3) 映射回原始输入张量 X 的位置 (j0, j1, j2, j3)j0 = (i0 - all_pads_0 + orig_dims_0) % orig_dims_0j1 = (i1 - all_pads_2 + orig_dims_1) % orig_dims_1j2 = (i2 - all_pads_4 + orig_dims_2) % orig_dims_2j3 = (i3 - all_pads_6 + orig_dims_3) % orig_dims_3load_idx = (orig_dims_3 * orig_dims_2 * orig_dims_1 * j0+ orig_dims_3 * orig_dims_2 * j1+ orig_dims_3 * j2+ j3)#计算输入张量 X 中的线性索引 load_idx，并通过 tl.load 加载数据（带越界检查 mask_x）mask_x = load_idx < X_len
#将加载的数据通过 tl.store 写入输出张量 Y 的对应位置（带越界检查 mask_y)x = tl.load(X + load_idx, mask=mask_x)tl.store(Y + i, x, mask=mask_y)@trtp.aot_impl("example::circ_pad_plugin")##@trtp.aot_impl标记，是TensorRT的预编译实现接口，负责将 Triton 内核编译为 CUDA 代码，并配置 TensorRT 运行时所需的参数。
def circ_pad_plugin_aot_impl(inp0: trtp.TensorDesc, pads: npt.NDArray[np.int32], outputs: Tuple[trtp.TensorDesc], tactic: int
) -> Tuple[Union[str, bytes], Union[str, bytes], trtp.KernelLaunchParams, trtp.SymExprs]:assert tactic == 0 #假设仅支持默认优化策略block_size = 256 #Triton内核的线程块大小type_str = "fp32" if inp0.dtype == trt.float32 else "fp16"#根据输入数据类型fp32/fp16生成类型字符串type_str用于内核签名### Triton内核编译#使用 triton.compiler.ASTSource 包装内核函数，指定其签名（输入输出类型）和编译时常量（BLOCK_SIZE）。src = triton.compiler.ASTSource(fn=circ_pad_kernel,signature=f"*{type_str},{','.join(['i32']*13)},*{type_str}",##签名格式constants={"BLOCK_SIZE": block_size,},)compiled_kernel = triton.compile(src)#调用triton.compile编译内核，生成包含PTX汇编代码和元数据的compiled_kernel.#### 内核启动该参数配置launch_params = trtp.KernelLaunchParams()N = inp0.ndimall_pads = np.zeros((N * 2,), dtype=np.int32)inp_dims = inp0.shape_exprout_dims = outputs[0].shape_exprfor i in range(np.size(pads) // 2):all_pads[N * 2 - 2 * i - 2] = pads[i * 2]all_pads[N * 2 - 2 * i - 1] = pads[i * 2 + 1]# grid dimslaunch_params.grid_x = trtp.cdiv(out_dims.numel(), block_size)##grid维度，根据输出张量的总元素数（out_dims.numel()）和 BLOCK_SIZE 计算（向上取整）。# block dimslaunch_params.block_x = compiled_kernel.metadata.num_warps * 32#线程块内的线程数 warpx32# shared memorylaunch_params.shared_mem = compiled_kernel.metadata.shared#内核使用的共享内存大小all_pads = all_pads.tolist()# Representing all int32 scalar inputs as symbolic expressions.# These inputs are either constants or derivatives of input/output shapes.# The symbolic expressions are resolved after the full shape context becomes available at runtime.# For the `circ_pad_kernel`, there are 13 such scalar extra-arguments, corresponding to the 13 arguments# between the mandatory input and output tensors in the triton kernel's function signature.
#将内核需要的 13 个标量参数（如填充量、输入 / 输出维度、总元素数）转换为 trtp.SymIntExprs（符号整数表达式）。
#这些符号参数在 TensorRT 运行时会根据实际输入 / 输出的形状解析为具体数值（例如，动态形状场景下，维度可能在运行时确定）。extra_args = trtp.SymIntExprs.from_tuple([trtp.SymInt32(e)for e in [all_pads[0],all_pads[2],all_pads[4],all_pads[6],inp_dims[0],inp_dims[1],inp_dims[2],inp_dims[3],out_dims[1],out_dims[2],out_dims[3],inp_dims.numel(),out_dims.numel(),]])return compiled_kernel.metadata.name, compiled_kernel.asm["ptx"], launch_params, extra_args ## 返回编译完的内核名称compiled_kernel.metadata.name## PTX代码：compiled_kernel.asm["ptx"]## 启动参数：launch_params## 符号参数：extra_args##以上内容供TensorRT在构建引擎时使用处理动态形状时，调用 tensorrt.plugin.aot_impl() 时可能无法知道具体的 I/O 维度。因此，内核启动参数和额外的内核参数以符号形式指定。
选择性能最佳的插件配置（自动调整autotuning）
如果插件能同时支持fp32和fp16，且不确定哪个性能更好
那么可以使用tensorrt.plugin.autotune()
@trtp.autotune("example::circ_pad_plugin")
def circ_pad_plugin_autotune(inp0: trtp.TensorDesc,pads: npt.NDArray[np.int32],outputs: Tuple[trtp.TensorDesc],
) -> List[trtp.AutoTuneCombination]:return [trtp.AutoTuneCombination("FP32|FP16, FP32|FP16", "LINEAR")]
Adding the plugin to a TensorRT network(将插件添加到TensorRT网络)
通过trt的python api 添加。api为：tensorrt.INetworkDefinition
Addding the plugin using TRT Python APIs
主要使用tensorrt.INetworkDefinition.add_plugin()可用于将插件添加到网络定义实例tensorrt.INetworkDefinition
input_tensor = network.add_input(name="x", dtype=trt.DataType.FLOAT, shape=x.shape)
plugin_layer = network.add_plugin(trt.plugin.op.example.circ_pad_plugin(input_tensor, pads = pads), aot = False)
注：
1、已经注册过的plugin可以在命名空间和tensorrt.plugin.op下找到
2、tensorrt.INetworkDefinition.add_plugin()采用boolean参数，可以使用aot参数选择究竟是AOT还是JIT实现。如果只有一个，那么可以不改
当Plugin同时定义了AOT和JIT实现后，可以通过全局设置来选择哪个实现
tensorrt.NetworkDefinitionCreationFlag.PREFER_AOT_PYTHON_PLUGINS
tensorrt.NetworkDefinitionCreationFlag.PREFER_JIT_PYTHON_PLUGINS
builder = trt.Builder(trt.Logger(trt.Logger.INFO))
# Always choose AOT implementations wherever possible
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.PREFER_AOT_PYTHON_PLUGINS))
...
# Now the `aot` argument can be omitted (and TRT will choose the AOT implementation)
plugin_layer = network.add_plugin(trt.plugin.op.example.circ_pad_plugin(input_tensor, pads = pads))
Loading an ONNX model with custom operator
使用自定义运算符加载ONNX模型
如果要通过TRT插件运行该插件，得保证：
1、ONNX 节点 op 的属性与您的插件名称完全相同。
2、该节点包含一个 string 属性，该属性以插件的命名空间命名 plugin_namespace 。
3、对于同时具有 AOT 和 JIT 实现的插件，必须制定实现方式
详情见此处
例如：使用ONNX Graphsurgeon，构建方式如下：
import onnx_graphsurgeon as gsvar_x = gs.Variable(name="x", shape=inp_shape, dtype=np.float32)
var_y = gs.Variable(name="y", dtype=np.float32)circ_pad_node = gs.Node(name="circ_pad_plugin",op="circ_pad_plugin",inputs=[var_x],outputs=[var_y],attrs={"pads": pads, "plugin_namespace": "example"},
)
高级用法
详情见此处
1、为具有数据相关输出形状的算子（如 non-zero）构建 shape 推导逻辑
2、TensorRT 插件中如何使用 .aliased() 方法实现就地计算操作，即让输出复用输入张量的内存，节省资源、提升性能。常用于加法、归一化等操作无需新内存开销的场景
3、如何在 TensorRT 插件中支持多个后端（backend）实现，并通过自动调优（autotune）选择性能最优的后端，即 自定义策略（tactic）机制。
TENSORRT PYTHON API REFERENCE
Foundational Types
DataType
Weights
略
Dims
略
IHostMemory
tensorrt.IHostMemory 是 TensorRT 用于返回如序列化模型等缓冲数据的对象，它支持 buffer 协议、由 TensorRT 自动管理生命周期，常通过 engine.serialize() 获得，并用于保存或传输模型数据。