层

有关更多信息，包括示例，请参阅 TensorRT 算子参考文档。

请注意，层权重属性可以表示为 NumPy 数组或 Weights 对象，具体取决于底层数据类型是否受 NumPy 支持。如果需要一致的类型，请从属性显式构造 Weights 对象

conv_layer = network.add_convolution_nd(...)
bias_weights = trt.Weights(conv_layer.bias)

PaddingMode

tensorrt.PaddingMode
枚举卷积、反卷积和池化层中可用的填充类型。

如果同时设置了 padding_mode 和 pre_padding，则填充模式优先。

EXPLICIT* 对应于显式填充。
SAME* 隐式计算填充，使输出维度与输入维度相同。对于卷积和池化，输出维度由 ceil(输入维度, 步幅) 确定。
CAFFE* 对应于对称填充。

成员

EXPLICIT_ROUND_DOWN : 使用显式填充，向下舍入输出大小

EXPLICIT_ROUND_UP : 使用显式填充，向上舍入输出大小

SAME_UPPER : 使用 SAME 填充，其中 pre_padding <= post_padding

SAME_LOWER : 使用 SAME 填充，其中 pre_padding >= post_padding

ICastLayer

class tensorrt.ICastLayer

表示类型转换功能的层。

此层将给定输入张量的元素转换为指定的数据类型，并返回转换类型中形状相同的输出张量。

支持除 FP8 之外的所有类型之间的转换。

变量:

to_type – DataType 输出张量的指定数据类型。

IConvolutionLayer

class tensorrt.IConvolutionLayer

INetworkDefinition 中的卷积层。

此层在 3 维或 4 维滤波器与 4 维或 5 维张量之间执行相关运算，以生成另一个 4 维或 5 维张量。

支持可选的偏置参数，该参数将每个通道的常量添加到输出中的每个值。

变量:

  • num_output_maps – int 卷积的输出特征图数量。

  • pre_padding – DimsHW 前填充。输入开始处将在高度和宽度方向上零填充此数量的元素。默认值：(0, 0)

  • post_padding – DimsHW 后填充。输入结束处将在高度和宽度方向上零填充此数量的元素。默认值：(0, 0)

  • padding_mode – PaddingMode 填充模式。如果同时设置了 IConvolutionLayer.padding_mode 和 IConvolutionLayer.pre_padding 或 IConvolutionLayer.post_padding 中的任一项，则填充模式优先。

  • num_groups – int 卷积的组数。输入张量通道被划分为这么多组，并且为每个组执行卷积，每个组使用一个滤波器。组卷积的结果被连接起来形成输出。 注意 当在 int8 模式下使用组时，组的大小（即通道计数除以组计数）对于输入和输出都必须是 4 的倍数。默认值：1。

  • kernel – Weights 卷积的核权重。权重被指定为 GKCRS 顺序的连续数组，其中 G 是组数，K 是输出特征图的数量，C 是输入通道的数量，R 和 S 是滤波器的高度和宽度。

  • bias – Weights 卷积的偏置权重。偏置是可选的。要省略偏置，请将其设置为空的 Weights 对象。偏置是按通道应用的，因此权重的数量（如果非零）必须等于输出特征图的数量。

  • kernel_size_nd – Dims 卷积的多维核大小。

  • stride_nd – Dims 卷积的多维步幅。默认值：(1, …, 1)

  • padding_nd – Dims 卷积的多维填充。输入将在每个维度上零填充此数量的元素。如果填充是不对称的，则此值对应于前填充。默认值：(0, …, 0)

  • dilation_nd – Dims 卷积的多维扩张。默认值：(1, …, 1)

IGridSampleLayer

tensorrt.InterpolationMode

插值的各种模式，用于调整大小和 grid_sample 层。

成员

NEAREST : 1D、2D 和 3D 最近邻插值。

LINEAR : 支持线性、双线性、三线性插值。

CUBIC : 支持双三次插值。

tensorrt.SampleMode

控制 ISliceLayer 和 IGridSample 如何处理越界坐标

成员

STRICT_BOUNDS : 当坐标越界时，报错失败。

WRAP : 坐标周期性地环绕。

CLAMP : 越界索引被钳制到边界

FILL : 当坐标越界时，使用填充输入值。

REFLECT : 坐标反射。

class tensorrt.IGridSampleLayer

INetworkDefinition 中的网格采样层。

此层使用输入张量和网格张量来生成插值输出张量。输入和网格张量必须是形状为 4 阶的张量。唯一支持的 SampleMode 是 trt.samplemode.CLAMP、trt.samplemode.FILL 和 trt.samplemode.REFLECT。

变量:

  • interpolation_mode – class:InterpolationMode 要使用的插值类型。默认为 LINEAR。

  • align_corners – class:bool 要使用的对齐模式。默认为 False。

  • sample_mode – SampleMode 要使用的采样模式。默认为 FILL。

IActivationLayer

tensorrt.ActivationType

要执行的激活类型。

成员

RELU : 修正线性激活

SIGMOID : Sigmoid 激活

TANH : 双曲正切激活

LEAKY_RELU : Leaky Relu 激活：f(x) = x 如果 x >= 0, f(x) = alpha * x 如果 x < 0

ELU : Elu 激活：f(x) = x 如果 x >= 0, f(x) = alpha * (exp(x) - 1) 如果 x < 0

SELU : Selu 激活：f(x) = beta * x 如果 x > 0, f(x) = beta * (alpha * exp(x) - alpha) 如果 x <= 0

SOFTSIGN : Softsign 激活：f(x) = x / (1 + abs(x))

SOFTPLUS : Softplus 激活：f(x) = alpha * log(exp(beta * x) + 1)

CLIP : Clip 激活：f(x) = max(alpha, min(beta, x))

HARD_SIGMOID : Hard sigmoid 激活：f(x) = max(0, min(1, alpha * x + beta))

SCALED_TANH : Scaled Tanh 激活：f(x) = alpha * tanh(beta * x)

THRESHOLDED_RELU : Thresholded Relu 激活：f(x) = x 如果 x > alpha, f(x) = 0 如果 x <= alpha

GELU_ERF : GELU erf 激活：0.5 * x * (1 + erf(sqrt(0.5) * x))

GELU_TANH : GELU tanh 激活：0.5 * x * (1 + tanh(sqrt(2/pi) * (0.044715F * pow(x, 3) + x)))

class tensorrt.IActivationLayer

INetworkDefinition 中的激活层。此层对其输入应用逐元素激活函数。输出的形状与输入相同。

变量:

  • type – ActivationType 要执行的激活类型。

  • alpha – float 一些参数化激活（LEAKY_RELU、ELU、SELU、SOFTPLUS、CLIP、HARD_SIGMOID、SCALED_TANH）使用的 alpha 参数。其他激活忽略此参数。

  • beta – float 一些参数化激活（SELU、SOFTPLUS、CLIP、HARD_SIGMOID、SCALED_TANH）使用的 beta 参数。其他激活忽略此参数。

IPoolingLayer

tensorrt.PoolingType

池化层中要执行的池化类型。

成员

MAX : 元素最大值

AVERAGE : 元素平均值。如果张量被填充，则计数包括填充

MAX_AVERAGE_BLEND : 最大池化和平均池化之间的混合：(1-blendFactor)*maxPool + blendFactor*avgPool

class tensorrt.IPoolingLayer

INetworkDefinition 中的池化层。该层在输入上的窗口内应用缩减操作。

变量:

  • type – PoolingType 要执行的池化类型。

  • pre_padding – DimsHW 前填充。输入开始处将在高度和宽度方向上零填充此数量的元素。默认值：(0, 0)

  • post_padding – DimsHW 后填充。输入结束处将在高度和宽度方向上零填充此数量的元素。默认值：(0, 0)

  • padding_mode – PaddingMode 填充模式。如果同时设置了 IPoolingLayer.padding_mode 和 IPoolingLayer.pre_padding 或 IPoolingLayer.post_padding 中的任一项，则填充模式优先。

  • blend_factor – float max_average_blend 模式的混合因子：\(max_average_blendPool = (1-blendFactor)*maxPool + blendFactor*avgPool\) 。 blend_factor 是 [0,1] 中的用户值，默认值为 0.0。此值仅适用于 PoolingType.MAX_AVERAGE_BLEND 模式。

  • average_count_excludes_padding – bool 平均池化是否使用窗口与未填充输入之间的重叠区域作为分母。如果未设置，则分母是池化窗口与填充输入之间的重叠。默认值：True

  • window_size_nd – Dims 池化的多维窗口大小。

  • stride_nd – Dims 池化的多维步幅。默认值：(1, …, 1)

  • padding_nd – Dims 池化的多维填充。默认值：(0, …, 0)

ILRNLayer

class tensorrt.ILRNLayer

INetworkDefinition 中的 LRN 层。输出大小与输入大小相同。

变量:

  • window_size – int LRN 窗口大小。窗口大小必须为奇数，且在 [1, 15] 范围内。

  • alpha – float LRN alpha 值。有效范围为 [-1e20, 1e20]。

  • beta – float LRN beta 值。有效范围为 [0.01, 1e5f]。

  • k – float LRN K 值。有效范围为 [1e-5, 1e10]。

IScaleLayer

tensorrt.ScaleMode

控制如何在 Scale 层中应用缩放。

成员

UNIFORM : 张量的所有元素具有相同的系数。

CHANNEL : 每通道系数。通道维度被假定为倒数第三个维度。

ELEMENTWISE : 逐元素系数。

class tensorrt.IScaleLayer

INetworkDefinition 中的 Scale 层。

此层对其输入应用逐元素计算

\(output = (input * scale + shift) ^ {power}\)

系数可以按每张量、每通道或每元素应用。

注意 如果权重的数量为 0，则 shift、power 和 scale 使用默认值。默认 shift 为 0，默认 power 为 1，默认 scale 为 1。

输出大小与输入大小相同。

注意 此层的输入张量需要至少有 3 个维度。

变量:

  • mode – ScaleMode 缩放模式。

  • shift – Weights shift 值。

  • scale – Weights scale 值。

  • power – Weights power 值。

  • channel_axis – int 通道轴。

ISoftMaxLayer

class tensorrt.ISoftMaxLayer

INetworkDefinition 中的 Softmax 层。

此层对其输入应用逐通道 softmax。

输出大小与输入大小相同。

变量:

axes – int 计算 softmax 的轴。目前，只能设置一个轴。

轴是通过将对应于该轴的位设置为 1 来指定的，作为一个位掩码。

例如，考虑一个 NCHW 张量作为输入（三个非批次维度）。

默认情况下，softmax 在轴上执行，该轴是轴数减三。如果非批次轴少于 3 个，则为 0。例如，如果输入为 NCHW，则默认轴为 C。如果输入为 NHW，则默认轴为 H。

位 0 对应于 N 维度布尔值。
位 1 对应于 C 维度布尔值。
位 2 对应于 H 维度布尔值。
位 3 对应于 W 维度布尔值。
默认情况下，softmax 在轴上执行，该轴是轴数减三。如果是 0
轴数少于 3 个。例如，如果输入为 NCHW，则默认轴为 C。如果输入
为 NHW，则默认轴为 N。

例如，要在 NPQRCHW 输入的 R 轴上执行 softmax，请设置位 3。

必须满足以下约束才能在 DLA 上执行此层

  • 轴必须是通道或空间维度之一。

  • 有两种受支持的输入大小类别

    • 非轴、非批次维度均为 1，轴维度最多为 8192。这是使用 softmax 的推荐情况，因为它最准确。

    • 至少一个非轴、非批次维度大于 1，轴维度最多为 1024。请注意，在这种情况下，当轴维度大小接近上限时，可能会存在一些近似误差。有关近似误差的更多详细信息，请参阅 TensorRT 开发者指南。

IConcatenationLayer

class tensorrt.IConcatenationLayer

INetworkDefinition 中的连接层。

输出通道大小是输入通道大小的总和。其他输出大小与其他输入大小相同，所有输入大小必须匹配。

变量:

axis – int 连接发生的轴。默认轴是张量维度数减三，如果张量的维度数少于三个，则为零。例如，对于维度为 NCHW 的张量，它是 C。

IDeconvolutionLayer

class tensorrt.IDeconvolutionLayer

INetworkDefinition 中的反卷积层。

变量:

  • num_output_maps – int 反卷积的输出特征图数量。

  • pre_padding – DimsHW 前填充。输入开始处将在高度和宽度方向上零填充此数量的元素。默认值：(0, 0)

  • post_padding – DimsHW 后填充。输入结束处将在高度和宽度方向上零填充此数量的元素。默认值：(0, 0)

  • padding_mode – PaddingMode 填充模式。如果同时设置了 IDeconvolutionLayer.padding_mode 和 IDeconvolutionLayer.pre_padding 或 IDeconvolutionLayer.post_padding 中的任一项，则填充模式优先。

  • num_groups – int 反卷积的组数。输入张量通道被划分为这么多组，并且为每个组执行反卷积，每个组使用一个滤波器。组反卷积的结果被连接起来形成输出。 注意 当在 int8 模式下使用组时，组的大小（即通道计数除以组计数）对于输入和输出都必须是 4 的倍数。默认值：1

  • kernel – Weights 反卷积的核权重。权重被指定为 CKRS 顺序的连续数组，其中 C 是输入通道的数量，K 是输出特征图的数量，R 和 S 是滤波器的高度和宽度。

  • bias – Weights 反卷积的偏置权重。偏置是可选的。要省略偏置，请将其设置为空的 Weights 对象。偏置是按特征图应用的，因此权重的数量（如果非零）必须等于输出特征图的数量。

  • kernel_size_nd – Dims 卷积的多维核大小。

  • stride_nd – Dims 反卷积的多维步幅。默认值：(1, …, 1)

  • padding_nd – Dims 反卷积的多维填充。输入将在每个维度上零填充此数量的元素。填充是对称的。默认值：(0, …, 0)

IElementWiseLayer

tensorrt.ElementWiseOperation

ElementWise 层可以执行的二元运算。

成员

SUM : 两个元素的和

PROD : 两个元素的积

MAX : 两个元素的最大值

MIN : 两个元素的最小值

SUB : 从第一个元素中减去第二个元素

DIV : 第一个元素除以第二个元素

POW : 第一个元素的第二个元素次幂

FLOOR_DIV : 第一个元素除以第二个元素的向下取整除法

AND : 两个元素的逻辑与

OR : 两个元素的逻辑或

XOR : 两个元素的逻辑异或

EQUAL : 检查两个元素是否相等

GREATER : 检查第一个张量中的元素是否大于第二个张量中的对应元素

LESS : 检查第一个张量中的元素是否小于第二个张量中的对应元素

class tensorrt.IElementWiseLayer

INetworkDefinition 中的逐元素层。

此层在两个张量的对应元素之间应用逐元素二元运算。

两个输入张量的输入维度必须相等，并且输出张量的大小与每个输入相同。

变量:

op – ElementWiseOperation 该层的二元运算。

IGatherLayer

class tensorrt.IGatherLayer

INetworkDefinition 中的 Gather 层。

变量:

  • axis – int 要在其上收集的非批次维度轴。轴必须小于数据输入中的非批次维度数。

  • num_elementwise_dims – int 要逐元素处理的索引张量的引导维度数。对于 GatherMode.DEFAULT，它可以是 0 或 1。对于 GatherMode::kND，它可以介于 0 和小于 rank(data) 的值之间。对于 GatherMode::kELEMENT，它必须为 0。

  • mode – GatherMode Gather 模式。

IPluginV2Layer

class tensorrt.IPluginV2Layer

INetworkDefinition 中的插件层。

变量:

plugin – IPluginV2 该层的插件。

IPluginV3Layer

class tensorrt.IPluginV3Layer

INetworkDefinition 中的插件层。

变量:

plugin – IPluginV3 该层的插件。

IUnaryLayer

tensorrt.UnaryOperation

Unary 层可能执行的一元运算。

成员

EXP : 指数运算

LOG : 对数运算 (底数为 e)

SQRT : 平方根

RECIP : 倒数

ABS : 绝对值

NEG : 取反

SIN : 正弦

COS : 余弦

TAN : 正切

SINH : 双曲正弦

COSH : 双曲余弦

ASIN : 反正弦

ACOS : 反余弦

ATAN : 反正切

ASINH : 反双曲正弦

ACOSH : 反双曲余弦

ATANH : 反双曲正切

CEIL : 向上取整

FLOOR : 向下取整

ERF : 高斯误差函数

NOT : 非

SIGN : 符号函数。如果输入 > 0，输出 1；如果输入 < 0，输出 -1；如果输入 == 0，输出 0。

ROUND : 对于浮点数据类型，四舍五入到最接近的偶数。

ISINF : 对于浮点数据类型，如果输入值等于 +/- 无穷大，则返回 true。

ISNAN : 对于浮点数据类型，如果输入值等于 NaN，则返回 true。

class tensorrt.IUnaryLayer

INetworkDefinition 中的 unary 层。

变量:

op – UnaryOperation 该层的 unary 运算。在 DLA 上运行此层时，仅支持 UnaryOperation.ABS。

IReduceLayer

tensorrt.ReduceOperation

Reduce 层可能执行的规约运算

成员

SUM

PROD

MAX

MIN

AVG

class tensorrt.IReduceLayer

INetworkDefinition 中的 reduce 层。

变量:

  • op – ReduceOperation 该层的规约运算。

  • axes – int 要在其上进行规约的轴。

  • keep_dims – bool 指定是否为该层保留规约后的维度。

IPaddingLayer

class tensorrt.IPaddingLayer

INetworkDefinition 中的 padding 层。

变量:

  • pre_padding_nd – Dims 在张量开始处应用的填充。负填充会导致按指定量修剪边缘。目前仅支持 2 个维度。

  • post_padding_nd – Dims 在张量末尾处应用的填充。负填充会导致按指定量修剪边缘。目前仅支持 2 个维度。

IParametricReLULayer

class tensorrt.IParametricReLULayer

INetworkDefinition 中的 parametric ReLU 层。

此层将参数化 ReLU 激活应用于输入张量（第一个输入），斜率取自斜率张量（第二个输入）。这可以看作是 leaky ReLU 运算，其中负斜率因元素而异（实际上可以学习）。

斜率张量必须可单向广播到输入张量：两个张量的秩必须相同，并且斜率张量的所有维度必须等于输入张量或为 1。输出张量的形状与输入张量相同。

ISelectLayer

class tensorrt.ISelectLayer

INetworkDefinition 中的 select 层。

此层实现元素级三元条件运算。在 condition 为 True 的位置，元素取自第一个输入；在 condition 为 False 的位置，元素取自第二个输入。

IShuffleLayer

class tensorrt.Permutation(*args, **kwargs)

排列的元素。排列的应用方式为 outputDimensionIndex = permutation[inputDimensionIndex]，因此要从 CHW 顺序排列到 HWC 顺序，所需的排列为 [1, 2, 0]，要从 HWC 排列到 CHW，所需的排列为 [2, 0, 1]。

它支持迭代和索引，并且可以隐式地与 Python 迭代器（如 tuple 或 list ）相互转换。因此，您可以将这些类代替 Permutation 使用。

重载函数。

  1. __init__(self: tensorrt.tensorrt.Permutation) -> None

  2. __init__(self: tensorrt.tensorrt.Permutation, arg0: List[int]) -> None

class tensorrt.IShuffleLayer

INetworkDefinition 中的 shuffle 层。

此类通过按顺序应用以下操作来混洗数据：转置操作、reshape 操作和第二次转置操作。输出的维度类型是 reshape 维度的维度类型。

变量:

  • first_transpose – Permutation 第一个转置操作应用的排列。默认值：单位排列

  • reshape_dims – Dims reshape 后的维度。两个特殊值可以用作维度。值 0 从输入复制相应的维度。此特殊值可以在维度中多次使用。如果 reshape 维度的数量少于输入，则通过对齐输入的最重要维度来解析 0。值 -1 通过查看输入和其余 reshape 维度来推断该特定维度。请注意，最多允许将一个维度指定为 -1。新维度的乘积必须等于旧维度的乘积。

  • second_transpose – Permutation 第二个转置操作应用的排列。默认值：单位排列

  • zero_is_placeholder – bool reshape 维度中 0 的含义。如果为 true，则 reshape 维度中的 0 表示从第一个输入张量复制相应的维度。如果为 false，则 reshape 维度中的 0 表示零长度维度。

set_input(self: tensorrt.tensorrt.ILayer, index: int, tensor: tensorrt.tensorrt.ITensor) → None

设置给定索引的输入张量。对于静态 shuffle 层，索引必须为 0。通过调用索引为 1 的 set_input() ，静态 shuffle 层将转换为动态 shuffle 层。动态 shuffle 层无法转换回静态 shuffle 层。

对于动态 shuffle 层，值 0 和 1 有效。动态情况下的索引如下

索引

描述

0

要混洗的数据或形状张量。

1

reshape 操作的维度，作为 1D int32 形状张量。

如果使用值 1 调用此函数，则 num_inputs 将从 1 更改为 2。

参数:

  • index – 输入张量的索引。

  • tensor – 输入张量。

ISliceLayer

class tensorrt.ISliceLayer

INetworkDefinition 中的 slice 层。

slice 层有两种变体：静态和动态。静态 slice 在层创建时通过 Dims 指定起始维度、大小和步长维度，并且可以使用 ISliceLayer 的 get/set 访问器函数。动态 slice 将起始维度、大小、步长或轴中的一个或多个指定为 ITensor`s， 通过使用 :func:`ILayer.set_input 添加第二个、第三个、第四个或第六个输入。如果输入缺失或为空，则使用相应的 Dims 。

应用程序可以根据大小或轴输入是否存在且非空来确定 ISliceLayer 是否具有动态输出形状。

slice 层为每个维度从输入张量中选择一个起始位置，并使用指定的步长跨输入张量将元素复制到输出张量。如果未通过 Dims 指定，则起始张量、大小张量和步长张量必须是 1 维整数类型形状张量。

在张量上使用 slice 的示例：input = {{0, 2, 4}, {1, 3, 5}} start = {1, 0} size = {1, 2} stride = {1, 2} output = {{1, 5}}

如果提供了轴，则起始维度、结束维度和步长维度必须与轴的长度相同，并指定要切片的维度子集。如果未提供轴，则起始维度、结束维度和步长维度必须与输入张量的秩长度相同。

使用指定轴在张量上使用 slice 的示例：input = {{0, 2, 4}, {1, 3, 5}} start = {1} size = {2} stride = {1} axes = {1} output = {{2, 4}, {3, 5}}

当 sampleMode 为 SampleMode.CLAMP 或 SampleMode.REFLECT 时，对于每个输入维度，如果其大小为 0，则相应的输出维度也必须为 0。

当 sampleMode 为 SampleMode.FILL 时，slice 层的第五个输入用于确定填充超出范围索引的值。在任何其他 sample mode 中指定第五个输入都是错误的。

如果满足以下条件，slice 层可以生成形状张量

  • start、size 和 stride 是构建时常量，可以是静态 Dims 或常量输入张量。

  • axes（如果提供）是构建时常量，可以是静态 Dims 或常量输入张量。

  • 输出张量中的元素数不超过 2 * Dims.MAX_DIMS 。

如果输出是形状张量，则输入张量是形状张量。

必须满足以下约束才能在 DLA 上执行此层：* start、size 和 stride 是构建时常量，可以是静态 Dims 或常量输入张量。 * axes（如果提供）是构建时常量，可以是静态 Dims 或常量输入张量。 * sampleMode 为 SampleMode.DEFAULT 、 SampleMode.WRAP 或 SampleMode.FILL 。 * 所有维度的步长均为 1。 * 不对第一个维度执行切片。 * 输入张量具有四个维度。 * 对于 SliceMode.FILL ，填充值输入是 IConstantLayer 的标量输出，值为 0，且未被任何其他层消耗。

变量:

  • start – Dims 起始偏移量。

  • shape – Dims 输出维度。

  • stride – Dims 切片步长。

  • mode – SampleMode 控制 ISliceLayer 如何处理越界坐标。

  • axes – Dims 起始维度、大小和步长维度对应的轴。

set_input(self: tensorrt.tensorrt.ILayer, index: int, tensor: tensorrt.tensorrt.ITensor) → None

设置给定索引的输入张量。对于静态 slice 层，索引必须为 0 或 4。通过调用索引介于 1 和 3 之间的 set_input() ，静态 slice 层将转换为动态 slice 层。动态 slice 层无法转换回静态 slice 层。

索引如下

索引

描述

0

要切片的数据或形状张量。

1

开始切片的起始张量，对于数据为 N 维，对于形状为 1 维。

2

结果切片的大小张量，对于数据为 N 维，对于形状为 1 维。

3

切片操作的步长，对于数据为 N 维，对于形状为 1 维。

4

用于 SampleMode.FILL 切片模式的值。不允许用于其他模式。

5

指示起始维度、大小和步长维度对应的轴的轴张量。必须是 1 维张量。

如果使用大于 0 的值调用此函数，则 num_inputs 将从 1 更改为 index + 1。

参数:

  • index – 输入张量的索引。

  • tensor – 输入张量。

IShapeLayer

class tensorrt.IShapeLayer

INetworkDefinition 中的 shape 层。用于获取张量的形状。此类将输出设置为一维张量，其维度与输入张量的维度相同。

例如，如果输入是具有维度 [2,3,5,7] 的四维张量（任何类型），则输出张量是长度为 4 的一维 int64 张量，其中包含序列 2、3、5、7。

ITopKLayer

tensorrt.TopKOperation

TopK 层可能执行的运算

成员

MAX : 元素最大值

MIN : 元素最小值

class tensorrt.ITopKLayer

INetworkDefinition 中的 TopK 层。

变量:

  • op – TopKOperation 该层的运算。

  • k – TopKOperation 该层的 k 值。目前仅支持高达 3840 的值。使用索引为 1 的 set_input() 方法传入动态 k 作为张量。

  • axes – TopKOperation 要在其上进行规约的轴。

set_input(self: tensorrt.tensorrt.ILayer, index: int, tensor: tensorrt.tensorrt.ITensor) → None

设置给定索引的输入张量。对于 TopK 层，索引必须为 0 或 1。

索引如下

索引

描述

0

输入数据张量。

1
一个标量 Int32 张量，包含与要检索的顶部元素数量相对应的正值

。大于 3840 的值将导致运行时错误。如果提供，这将覆盖计算中的静态 k 值。
参数:

  • index – 输入张量的索引。

  • tensor – 输入张量。

IMatrixMultiplyLayer

tensorrt.MatrixOperation

Matrix 层可能执行的矩阵运算

成员

NONE

TRANSPOSE : 转置每个矩阵

VECTOR : 将操作数视为向量集合

class tensorrt.IMatrixMultiplyLayer

INetworkDefinition 中的矩阵乘法层。

设 A 为 op(getInput(0))，B 为 op(getInput(1))，其中 op(x) 表示相应的 MatrixOperation。

当 A 和 B 是矩阵或向量时，计算内积 A * B

矩阵 * 矩阵 -> 矩阵
矩阵 * 向量 -> 向量
向量 * 矩阵 -> 向量
向量 * 向量 -> 标量

更高秩的输入被视为矩阵或向量的集合。输出将是相应的矩阵、向量或标量的集合。

变量:

IRaggedSoftMaxLayer

class tensorrt.IRaggedSoftMaxLayer

INetworkDefinition 中的 ragged softmax 层。

此层采用 ZxS 输入张量和额外的 Zx1 边界张量，其中包含 Z 序列的长度。

此层计算每个 Z 序列的 softmax。

输出张量的大小与输入张量相同。

IIdentityLayer

class tensorrt.IIdentityLayer

一个表示恒等函数的层。

如果显式指定了张量精度，则可以使用它从一种精度转换为另一种精度。

除了相同类型之间的转换（例如，float32 -> float32），唯一有效的转换是

(float32 | float16 | int32 | bool) -> (float32 | float16 | int32 | bool)

(float32 | float16) -> uint8

uint8 -> (float32 | float16)

IConstantLayer

class tensorrt.IConstantLayer

INetworkDefinition 中的一个常量层。

注意：此层不支持布尔和 uint8 类型。

变量:

  • weights – Weights 该层的权重。

  • shape – Dims 该层的形状。

IResizeLayer

class tensorrt.IResizeLayer

INetworkDefinition 中的一个调整大小层。

调整大小层可用于调整 N 维张量的大小。

调整大小层目前支持以下配置

  • InterpolationMode.NEAREST - 调整 N 维张量最内层 m 个维度的大小，其中 0 < m <= min(3, N) 且 N > 0。

  • InterpolationMode.LINEAR - 调整 N 维张量最内层 m 个维度的大小，其中 0 < m <= min(3, N) 且 N > 0。

  • InterpolationMode.CUBIC - 调整 N 维张量最内层 2 个维度的大小，N >= 2。

默认调整大小模式为 InterpolationMode.NEAREST。

调整大小层提供两种调整张量维度的方法

  • 直接设置输出维度。静态和动态调整大小层都可以这样做。

    静态调整大小层要求在构建时知道输出维度。动态调整大小层要求将输出维度设置为输入张量之一。

  • 设置调整大小比例。每个输出维度计算为 floor(输入维度 * 比例)。

    只有静态调整大小层允许设置在构建时已知的比例。

如果在 DLA 上执行此层，则支持以下参数组合

  • 在 NEAREST 模式下

    • (ResizeCoordinateTransformation.ASYMMETRIC，ResizeSelector.FORMULA，ResizeRoundMode.FLOOR)

    • (ResizeCoordinateTransformation.HALF_PIXEL，ResizeSelector.FORMULA，ResizeRoundMode.HALF_DOWN)

    • (ResizeCoordinateTransformation.HALF_PIXEL，ResizeSelector.FORMULA，ResizeRoundMode.HALF_UP)

  • 在 LINEAR 和 CUBIC 模式下

    • (ResizeCoordinateTransformation.HALF_PIXEL，ResizeSelector.FORMULA)

    • (ResizeCoordinateTransformation.HALF_PIXEL，ResizeSelector.UPPER)

变量:

  • shape – Dims 输出维度。必须等于输入维度大小。

  • scales – List[float] 调整大小比例列表。如果在 DLA 上执行此层，则有三个限制：1. len(scales) 必须正好为 4。2. 比例中的前两个元素必须正好为 1（对于未更改的批次和通道维度）。3. 最后两个元素在比例中，分别表示高度和宽度维度上的比例值，对于 NEAREST 模式需要是 [1, 32] 范围内的整数值，对于 LINEAR 模式需要是 [1, 4] 范围内的整数值。DLA 支持的比例示例：[1, 1, 2, 2]。

  • resize_mode – InterpolationMode 调整大小模式可以是 Linear、Cubic 或 Nearest。

  • coordinate_transformation – ResizeCoordinateTransformationDoc 支持的调整大小坐标变换模式为 ALIGN_CORNERS、ASYMMETRIC 和 HALF_PIXEL。

  • selector_for_single_pixel – ResizeSelector 支持的调整大小选择器模式为 FORMULA 和 UPPER。

  • nearest_rounding – ResizeRoundMode 支持的调整大小舍入模式为 HALF_UP、HALF_DOWN、FLOOR 和 CEIL。

  • exclude_outside – int 如果设置为 1，则输入张量外部采样位置的权重将设置为 0，并且权重将被重新归一化，使其总和为 1.0。

  • cubic_coeff – float 立方插值中使用的系数 'a'。

set_input(self: tensorrt.tensorrt.ILayer, index: int, tensor: tensorrt.tensorrt.ITensor) → None

为给定索引设置输入张量。

如果 index == 1 且 num_inputs == 1，则 num_inputs 更改为 2。一旦设置了这样的附加输入，调整大小层将在动态模式下工作。当 index == 1 且 num_inputs == 1 时，输出维度将从输入张量中使用，覆盖由 shape 提供的维度。

参数:

  • index – 输入张量的索引。

  • tensor – 输入张量。

ILoop

class tensorrt.ILoop

用于创建循环子图的助手。

变量:

name – 循环的名称。该名称用于错误诊断。

add_iterator(self: tensorrt.tensorrt.ILoop, tensor: tensorrt.tensorrt.ITensor, axis: int = 0, reverse: bool = False) → tensorrt.tensorrt.IIteratorLayer

返回通过循环迭代下标化张量的层。

对于 reverse=false，这等效于 add_gather(tensor, I, 0)，其中 I 是包含循环迭代次数的标量张量。对于 reverse=true，这等效于 add_gather(tensor, M-1-I, 0)，其中 M 是从 COUNT 类型的 TripLimits 计算出的行程计数。

参数:

  • tensor – 要迭代的张量。

  • axis – 要沿其迭代的轴。

  • reverse – 是否反向迭代。

返回值:

IIteratorLayer ，如果无法创建，则为 None 。

add_loop_output(self: tensorrt.tensorrt.ILoop, tensor: tensorrt.tensorrt.ITensor, kind: tensorrt.tensorrt.LoopOutput, axis: int = 0) → tensorrt.tensorrt.ILoopOutputLayer

为此循环创建基于给定张量的输出。

如果 kind 为 CONCATENATE 或 REVERSE，则必须通过方法 ILoopOutputLayer.set_input() 添加指定连接维度的第二个输入。

参数:

  • kind – 循环输出的类型。请参阅 LoopOutput

  • axis – 用于连接的轴（如果使用 CONCATENATE 或 REVERSE 类型的 kind）。

返回值:

添加的 ILoopOutputLayer ，如果无法创建，则为 None 。

add_recurrence(self: tensorrt.tensorrt.ILoop, initial_value: tensorrt.tensorrt.ITensor) → tensorrt.tensorrt.IRecurrenceLayer

为此循环创建以 initial_value 作为其第一个输入的循环层。

参数:

initial_value – 循环层的初始值。

返回值:

添加的 IRecurrenceLayer ，如果无法创建，则为 None 。

add_trip_limit(self: tensorrt.tensorrt.ILoop, tensor: tensorrt.tensorrt.ITensor, kind: tensorrt.tensorrt.TripLimit) → tensorrt.tensorrt.ITripLimitLayer

添加基于给定张量的行程计数限制器。

一个循环最多可以有一个 COUNT 和一个 WHILE 限制器。当同时存在行程限制时，循环在达到计数或条件为假时退出。不添加至少一个行程限制器是错误的。

对于 WHILE，输入张量必须是不包含 ITripLimitLayer 、 IIteratorLayer 或 ILoopOutputLayer 的层的子图的输出。子图中的任何 IRecurrenceLayer 都必须与 ITripLimitLayer 属于同一循环。此规则的一个简单示例是 WHILE 的输入是同一循环的 IRecurrenceLayer 的输出。

参数:

  • tensor – 输入张量。必须在循环开始之前可用。

  • kind – 行程限制的类型。请参阅 TripLimit

返回值:

添加的 ITripLimitLayer ，如果无法创建，则为 None 。

ILoopBoundaryLayer

class tensorrt.ILoopBoundaryLayer
变量:

loop – 与此边界层关联的 ILoop 。

ITripLimitLayer

tensorrt.TripLimit

描述行程限制的类型。

成员

COUNT : Tensor 是类型为 int32 的标量，其中包含行程计数。

WHILE : Tensor 是类型为 bool 的标量。当其值为 false 时，循环终止。

class tensorrt.ITripLimitLayer
变量:

kind – 行程限制的类型。请参阅 TripLimit

IRecurrenceLayer

class tensorrt.IRecurrenceLayer
set_input(self: tensorrt.tensorrt.ILayer, index: int, tensor: tensorrt.tensorrt.ITensor) → None

设置第一个或第二个输入。如果 index==1 且输入数为 1，则会追加输入。第一个输入指定初始输出值，并且必须来自循环外部。第二个输入指定下一个输出值，并且必须来自循环内部。两个输入必须具有相同的维度。

参数:

  • index – 要设置的输入的索引。

  • tensor – 输入张量。

IIteratorLayer

class tensorrt.IIteratorLayer
变量:

  • axis – 要迭代的轴

  • reverse – 对于 reverse=false，该层等效于 add_gather(tensor, I, 0)，其中 I 是包含循环迭代次数的标量张量。对于 reverse=true，该层等效于 add_gather(tensor, M-1-I, 0)，其中 M 是从 COUNT 类型的 TripLimits 计算出的行程计数。默认值为 reverse=false。

ILoopOutputLayer

tensorrt.LoopOutput

描述循环输出的类型。

成员

LAST_VALUE : 输出值是最后一次迭代的张量值。

CONCATENATE : 输出值是每次迭代的张量值的连接，按正向顺序排列。

REVERSE : 输出值是每次迭代的张量值的连接，按反向顺序排列。

class tensorrt.ILoopOutputLayer

ILoopOutputLayer 是从循环获取输出的唯一方法。

第一个输入张量必须在循环内定义；输出张量在循环外部。第二个输入张量（如果存在）必须在循环外部定义。

如果 kind 为 LAST_VALUE，则必须提供单个输入。

如果 kind 为 CONCATENATE 或 REVERSE，则必须提供第二个输入。第二个输入必须是在循环开始之前定义的标量“形状张量”，用于指定输出的连接长度。

输出张量比输入张量多 j 个维度，其中如果 kind 为 LAST_VALUE 则 j == 0，如果 kind 为 CONCATENATE 或 REVERSE 则 j == 1。

变量:

  • axis – 连接轴。如果 kind 为 LAST_VALUE，则忽略。例如，如果输入张量的维度为 [b,c,d]，并且 kind 为 CONCATENATE，则输出具有四个维度。设 a 为第二个输入的值。axis=0 会导致输出具有维度 [a,b,c,d]。axis=1 会导致输出具有维度 [b,a,c,d]。axis=2 会导致输出具有维度 [b,c,a,d]。axis=3 会导致输出具有维度 [b,c,d,a]。默认为 axis 为 0。

  • kind – 循环输出的类型。请参阅 LoopOutput

set_input(self: tensorrt.tensorrt.ILayer, index: int, tensor: tensorrt.tensorrt.ITensor) → None

类似于 ILayer.set_input()，但如果 index==1，num_inputs`==1, in which case :attr:`num_inputs 更改为 2，则额外有效。

IFillLayer

tensorrt.FillOperation

填充层可以执行的张量填充操作。

成员

LINSPACE : 在指定区间内生成均匀间隔的数字

RANDOM_UNIFORM : 生成一个张量，其随机值来自均匀分布

RANDOM_NORMAL : 生成一个张量，其随机值来自正态分布

class tensorrt.IFillLayer

INetworkDefinition 中的一个填充层。

输出张量的数据类型可以通过 to_type 指定。每个填充操作支持的输出类型如下。

操作

to_type

kLINSPACE

int32, int64, float32

kRANDOM_UNIFORM

float16, float32

kRANDOM_NORMAL

float16, float32
变量:

to_type – DataType 输出张量的指定数据类型。默认为 tensorrt.float32。

is_alpha_beta_int64(self: tensorrt.tensorrt.IFillLayer) → bool
set_input(self: tensorrt.tensorrt.ILayer, index: int, tensor: tensorrt.tensorrt.ITensor) → None

用特定的张量替换此层的输入。

索引

kLINSPACE 的描述

0

Shape 张量，表示输出张量的维度。

1

Start，标量，表示起始值。

2

Delta，1D 张量，长度等于 shape 张量的 nbDims，表示每个维度的 delta 值。

索引

kRANDOM_UNIFORM 的描述

0

Shape 张量，表示输出张量的维度。

1

Minimum，标量，表示最小随机值。

2

Maximum，标量，表示最大随机值。

索引

kRANDOM_NORMAL 的描述

0

Shape 张量，表示输出张量的维度。

1

Mean，标量，表示正态分布的均值。

2

Scale，标量，表示正态分布的标准差。
参数:

  • index – 要修改的输入的索引。

  • tensor – 输入张量。

IQuantizeLayer

class tensorrt.IQuantizeLayer

INetworkDefinition 中的量化层，属于 INetworkDefinition 。

此层接受浮点数据输入张量，并使用 scale 和 zeroPt 输入来

根据以下公式将数据量化为 8 位有符号整数

\(output = clamp(round(input / scale) + zeroPt)\)

舍入类型为四舍五入到最接近的偶数（https://en.wikipedia.org/wiki/Rounding#Round_half_to_even）。

Clamp 范围为 [-128, 127]。

第一个输入（索引 0）是要量化的张量。第二个（索引 1）和第三个（索引 2）分别是 scale 和零点。scale 和 zeroPt 必须都是标量或 1D 张量。

zeroPt 张量是可选的，如果未设置，则假定为零。其数据类型必须为 tensorrt.int8。zeroPt 必须仅包含零值系数，因为仅支持对称量化。scale 值必须是用于逐张量量化的标量，或用于逐轴量化的 1D 张量。1-D scale 张量的大小必须与量化轴的大小匹配。scale 的大小必须与 zeroPt 的大小匹配。

以 scale 张量结尾的子图必须是构建时常量。相同的限制适用于 zeroPt。输出类型（如果受约束）必须约束为 tensorrt.int8 或 tensorrt.fp8。输入类型（如果受约束）必须约束为 tensorrt.float32、tensorrt.float16 或 tensorrt.bfloat16。输出大小与输入大小相同。

IQuantizeLayer 支持 tensorrt.float32、tensorrt.float16 和 tensorrt.bfloat16 精度，并且在实例化期间将默认为 tensorrt.float32 精度。IQuantizeLayer 支持 tensorrt.int8 和 tensorrt.float8 输出。

变量:

  • axis – int 量化发生的轴。量化轴参考输入张量的维度。

  • to_type – DataType 指定的输出张量数据类型。必须为 tensorrt.int8 或 tensorrt.float8。

IDequantizeLayer

class tensorrt.IDequantizeLayer

INetworkDefinition 中的反量化层，属于 INetworkDefinition 。

此层接受有符号 8 位整数输入张量，并使用配置的 scale 和 zeroPt 输入根据以下公式反量化输入：\(output = (input - zeroPt) * scale\)

第一个输入（索引 0）是要量化的张量。第二个（索引 1）和第三个（索引 2）分别是 scale 和零点。scale 和 zeroPt 必须都是标量或 1D 张量。

zeroPt 张量是可选的，如果未设置，则假定为零。其数据类型必须为 tensorrt.int8。zeroPt 必须仅包含零值系数，因为仅支持对称量化。scale 值必须是用于逐张量量化的标量，或用于逐轴量化的 1D 张量。1-D scale 张量的大小必须与量化轴的大小匹配。scale 的大小必须与 zeroPt 的大小匹配。

以 scale 张量结尾的子图必须是构建时常量。相同的限制适用于 zeroPt。输出类型（如果受约束）必须约束为 tensorrt.int8 或 tensorrt.fp8。输入类型（如果受约束）必须约束为 tensorrt.float32、tensorrt.float16 或 tensorrt.bfloat16。输出大小与输入大小相同。

IDequantizeLayer 支持 tensorrt.int8 和 tensorrt.float8 精度，并且在实例化期间将默认为 tensorrt.int8 精度。IDequantizeLayer 支持 tensorrt.float32、tensorrt.float16 和 tensorrt.bfloat16 输出。

变量:

  • axis – int 反量化发生的轴。反量化轴参考输入张量的维度。

  • to_type – DataType 指定的输出张量数据类型。必须为 tensorrt.float32 或 tensorrt.float16。

IDynamicQuantizeLayer

class tensorrt.IDynamicQuantizeLayer

INetworkDefinition 中的动态量化层，属于 INetworkDefinition 。

此层对其输入张量执行动态块量化，并输出量化数据和计算出的块缩放因子。块大小目前限制为 16，且分块轴的大小必须可被 16 整除。

第一个输入（索引 0）是要量化的张量。其数据类型必须为 DataType::kFLOAT、DataType::kHALF 或 DataType::kBF16 之一。目前仅支持 2D 和 3D 输入。

第二个输入（索引 1）是双量化缩放因子。它是用于量化计算出的块缩放因子的标量缩放因子。

变量:

  • axis – int 被切分成块的轴。该轴必须是最后一个维度或倒数第二个维度。

  • block_size – int 使用共享缩放因子量化的元素数量。目前仅支持 16 个元素的块。

  • output_type – DataType 量化输出张量的数据类型，必须为 DataType::kFP4。

  • scale_type – DataType 用于量化输入数据的缩放因子数据类型，必须为 DataType::kFP8。

IScatterLayer

class tensorrt.IScatterLayer

INetworkDefinition 中的 Scatter 层，属于 INetworkDefinition。 :ivar axis: 在使用 Scatter Element 模式时进行 scatter 操作的轴（在 ND 模式下忽略） :ivar mode: ScatterMode scatter 的操作模式。

IIfConditional

class tensorrt.IIfConditional

用于构建条件执行子图的辅助类。

If-conditional 根据以下伪代码有条件地执行（惰性求值）网络的一部分

If condition is true Then:
    output = trueSubgraph(trueInputs);
Else:
    output = falseSubgraph(falseInputs);
Emit output

Condition 是 0D 布尔张量（表示标量）。trueSubgraph 表示当条件评估为 True 时执行的网络子图。falseSubgraph 表示当条件评估为 False 时执行的网络子图。

以下约束适用于 If-conditional： - trueSubgraph 和 falseSubgraph 都必须已定义。 - 两个子图中的输出张量数量相同。 - 来自 true/false 子图的每个输出张量的类型和形状都相同，但如果条件是构建时常量，则允许形状不同。

add_input(self: tensorrt.tensorrt.IIfConditional, input: tensorrt.tensorrt.ITensor) → tensorrt.tensorrt.IIfConditionalInputLayer

基于给定的张量，为此 if-conditional 创建一个输入。

参数:

input – 条件的输入，可供条件的子图之一或两者使用。

add_output(self: tensorrt.tensorrt.IIfConditional, true_subgraph_output: tensorrt.tensorrt.ITensor, false_subgraph_output: tensorrt.tensorrt.ITensor) → tensorrt.tensorrt.IIfConditionalOutputLayer

基于给定的张量，为此 if-conditional 创建一个输出。

if-conditional 的每个输出层表示 if-conditional 的 true-subgraph 或 false-subgraph 的单个输出，具体取决于执行哪个子图。

参数:

  • true_subgraph_output – 当此条件的条件输入评估为 true 时执行的子图的输出。

  • false_subgraph_output – 当此条件的条件输入评估为 false 时执行的子图的输出。

返回值:

IIfConditionalOutputLayer ，如果无法创建，则为 None。

set_condition(self: tensorrt.tensorrt.IIfConditional, condition: tensorrt.tensorrt.ITensor) → tensorrt.tensorrt.IConditionLayer

为此 If-Conditional 构造设置条件张量。

condition 张量必须是类型为 bool 的 0D 数据张量（标量）。

参数:

condition – 确定要执行哪个子图的条件张量。

返回值:

IConditionLayer ，如果无法创建，则为 None。

IConditionLayer

class tensorrt.IConditionLayer

描述 if-conditional 的布尔条件。

IIfConditionalOutputLayer

class tensorrt.IIfConditionalOutputLayer

描述 if-conditional 输出的类型。

IIfConditionalInputLayer

class tensorrt.IIfConditionalInputLayer

描述 if-conditional 输入的类型。

IEinsumLayer

class tensorrt.IEinsumLayer

INetworkDefinition 中的 Einsum 层，属于 INetworkDefinition 。

此层根据爱因斯坦求和约定，沿方程参数指定的维度对输入的元素执行求和。该层可以有一个或多个秩 >= 0 的输入。所有输入的数据类型必须相同。此层支持除 bool 之外的所有 TensorRT 数据类型。有一个与输入张量类型相同的输出张量。输出张量的形状由方程确定。

方程为每个输入中的每个维度指定 ASCII 小写字母，维度顺序与字母顺序相同，每个输入之间用逗号分隔。用相同下标标记的维度必须匹配或可广播。一个输入中重复的下标标签取对角线。跨多个输入重复标签意味着这些轴将被相乘。从输出中省略标签意味着将对沿这些轴的值求和。在隐式模式下，表达式中出现一次的索引将按字母顺序递增排列在输出中。在显式模式下，可以通过指定输出下标标签来控制输出，方法是添加箭头（'->'）后跟输出的下标。例如，“ij,jk->ik”等效于“ij,jk”。省略号（'...'）可以代替下标来广播维度。有关方程语法的更多详细信息，请参阅 TensorRT 开发者指南。

许多常见操作可以使用 Einsum 方程表示。例如： 矩阵转置：ij->ji 求和：ij-> 矩阵-矩阵乘法：ik,kj->ij 点积：i,i-> 矩阵-向量乘法：ik,k->i 批矩阵乘法：ijk,ikl->ijl 批对角线：…ii->…i

请注意，TensorRT 不支持省略号或对角线操作。

变量:

equation – str 层的 Einsum 方程。方程是以逗号分隔的下标标签列表，其中每个标签都引用相应张量的维度。

IAssertionLayer

class tensorrt.IAssertionLayer

INetworkDefinition 中的断言层，属于 INetworkDefinition 。

此层实现断言。输入必须是布尔形状张量。如果它的任何元素为 False，则会发生构建时或运行时错误。断言输入维度的相等性可能有助于优化器。

变量:

message – string 如果断言失败要打印的消息。

IOneHotLayer

class tensorrt.IOneHotLayer

网络定义中的 OneHot 层。

OneHot 层有三个输入张量：Indices、Values 和 Depth，一个输出张量 Output，以及一个 axis 属性。 :ivar indices: 是一个 Int32 张量，用于确定 Output 中哪些位置应设置为 on_value。 :ivar values: 是一个由 [off_value, on_value] 组成的双元素（秩=1）张量 :ivar depth: 是一个秩为 0 的 Int32 形状张量，其中包含 one-hot 编码的深度（类别数）。depth 张量必须是构建时常量，并且其值应为正数。 :returns: 秩 = rank(indices)+1 的张量，其中添加的维度包含 one-hot 编码。 :param axis: 指定 one-hot 编码添加到输出的哪个维度。

Output 的数据类型应等于 Values 数据类型。输出的计算方法是将 off_values 复制到所有输出元素，然后在 indices 张量指定的索引上设置 on_value。

当 axis = 0 时：对于所有 i、j、k，output[indices[i, j, k], i, j, k] = on_value，否则为 off_value。

当 axis = -1 时：对于所有 i、j、k，output[i, j, k, indices[i, j, k]] = on_value，否则为 off_value。

INonZeroLayer

class tensorrt.INonZeroLayer

INetworkDefinition 中的 NonZero 层，属于 INetworkDefinition 。

计算输入张量中值不为零的索引。返回的索引按行优先顺序排列。

输出形状始终为 {D, C}，其中 D 是输入的维度数， C 是非零值的数量。

INMSLayer

class tensorrt.INMSLayer

INetworkDefinition 中的非极大值抑制层，属于 INetworkDefinition 。

Boxes：层的输入 boxes 张量。此张量包含输入边界框。它是 float32 或 float16 类型的线性张量。如果框是按类别划分的，则其形状为 [batchSize, numInputBoundingBoxes, numClasses, 4]，如果相同的框要用于每个类别，则形状为 [batchSize, numInputBoundingBoxes, 4]。

Scores：层的输入 scores 张量。此张量包含每个框的分数。它是与 boxes 张量类型相同的线性张量。其形状为 [batchSize, numInputBoundingBoxes, numClasses]。

MaxOutputBoxesPerClass：层的输入 maxOutputBoxesPerClass 张量。此张量包含每个批次项、每个类别的最大输出框数。它是 int32 类型的标量（0D 张量）。

IoUThreshold 是所选框的最大 IoU。它是 float32 类型的标量（0D 张量），范围为 [0.0, 1.0]。它是可选输入，默认值为 0.0。使用 set_input() 添加此可选张量。

ScoreThreshold 是框分数要超过才能被选中的值。它是 float32 类型的标量（0D 张量）。它是可选输入，默认值为 0.0。使用 set_input() 添加此可选张量。

SelectedIndices 输出张量包含所选框的索引。它是 int32 类型的线性张量。其形状为 [NumOutputBoxes, 3]。] 每行包含一个 (batchIndex, classIndex, boxIndex) 元组。输出框按 batchIndex 递增的顺序排序，然后在每个 batchIndex 内按分数递减的顺序排序。对于每个 batchIndex，具有相同分数的输出框的顺序未指定。如果 MaxOutputBoxesPerClass 是常量输入，则最大输出框数为 batchSize * numClasses * min(numInputBoundingBoxes, MaxOutputBoxesPerClass)。否则，最大输出框数为 batchSize * numClasses * numInputBoundingBoxes。最大输出框数用于确定此输出张量分配内存的上限。

NumOutputBoxes 输出张量包含 selectedIndices 中的输出框数。它是 int32 类型的标量（0D 张量）。

NMS 算法按分数降序迭代一组边界框及其置信度分数。如果框的分数高于给定阈值，并且其与先前选择的框的交并比 (IoU) 小于或等于给定阈值，则选择该框。此层实现每个批次项和每个类别的 NMS。

对于每个批次项，具有相同分数的候选边界框的顺序未指定。

变量:

  • bounding_box_format – BoundingBoxFormat 层使用的边界框格式。默认为 CORNER_PAIRS。

  • topk_box_limit – int 筛选出的框中用于选择的最大数量。SM 5.3 和 6.2 设备默认为 2000，其他设备默认为 5000。TopK 框限制必须小于或等于 {SM 5.3 和 6.2 设备为 2000，其他设备为 5000}。

set_input(self: tensorrt.tensorrt.ILayer, index: int, tensor: tensorrt.tensorrt.ITensor) → None

为给定索引设置输入张量。索引如下所示

索引

描述

0

必需的 Boxes 张量。

1

必需的 Scores 张量。

2

必需的 MaxOutputBoxesPerClass 张量。

3

可选的 IoUThreshold 张量。

4

可选的 ScoreThreshold 张量。

如果为大于或等于 num_inputs 的索引调用此函数，则之后 num_inputs 返回 index + 1，并且任何缺失的中间输入都将设置为 null。请注意，只有可选输入可以缺失。

参数:

  • index – 输入张量的索引。

  • tensor – 输入张量。

IReverseSequenceLayer

class tensorrt.IReverseSequenceLayer

INetworkDefinition 中的 ReverseSequence 层，属于 INetworkDefinition 。

此层执行按批次反转，它沿轴 batch_axis 对输入张量进行切片。对于第 i 个切片，该操作沿 sequence_axis 反转前 N 个元素（由 sequence_lens 中相应的第 i 个值指定），并保持其余元素不变。输出张量将具有与输入张量相同的形状。

变量:

  • batch_axis – int 批次轴。默认值：1。

  • sequence_axis – int 序列轴。默认值：0。

INormalizationLayer

class tensorrt.INormalizationLayer

INetworkDefinition 中的归一化层，属于 INetworkDefinition 。

归一化层执行以下操作

X - 输入张量 Y - 输出张量 S - 比例张量 B - 偏置张量

Y = (X - Mean(X, axes)) / Sqrt(Variance(X) + epsilon) * S + B

其中 Mean(X, axes) 是对一组轴的缩减，而 Variance(X) = Mean((X - Mean(X, axes)) ^ 2, axes)。

变量:

  • epsilon – float 用于归一化计算的 epsilon 值。默认值：1e-5F。

  • axes – int 归一化计算的缩减轴。

  • num_groups – int 将通道分成组进行归一化计算的组数。默认值：1。

  • compute_precision – DataType 此层计算精度所使用的数据类型。默认情况下，即使在混合精度模式下，无论设置任何构建器标志，TensorRT 都会以 DataType.kFLOAT32 数据类型运行归一化计算，以避免溢出错误。ILayer.precision 和 ILayer.set_output_type 仍然可以设置以控制此层的输入和输出类型。此成员仅对 DataType.kFLOAT32 和 DataType.kHALF 有效。默认值：Datatype.FLOAT。

ISqueezeLayer

class tensorrt.ISqueezeLayer

INetworkDefinition 中的 Squeeze 层。

此层表示 squeeze 操作，移除输入张量在一组轴上的单位维度。

轴必须可解析为常量 Int32 或 Int64 1D 形状张量。轴中的值必须是唯一的，并且在 [-r, r-1] 范围内，其中 r 是输入张量的秩。对于每个轴值，输入张量中的相应维度必须为 1。

set_input(self: tensorrt.tensorrt.ILayer, index: int, tensor: tensorrt.tensorrt.ITensor) → None

为给定的索引设置输入张量。对于 Squeeze 层，索引必须为 0 或 1。

索引如下

索引

描述

0

输入数据张量。

1

要移除的轴。必须可解析为常量 Int32 或 Int64 1D 形状张量。
参数:

  • index – 输入张量的索引。

  • tensor – 输入张量。

IUnsqueezeLayer

class tensorrt.IUnsqueezeLayer

INetworkDefinition 中的 Unsqueeze 层。

此层表示 unsqueeze 操作，通过在输出的指定轴上插入单位长度维度来重塑输入张量。

轴必须可解析为常量 Int32 或 Int64 形状张量。轴中的值必须是唯一的，并且在 [-r_final, r_final-1] 范围内，其中 r_final 是 rank(input) 和 len(axes) 的和。

r_final 必须小于 Dims.MAX_DIMS。

set_input(self: tensorrt.tensorrt.ILayer, index: int, tensor: tensorrt.tensorrt.ITensor) → None

为给定的索引设置输入张量。对于 Unsqueeze 层，索引必须为 0 或 1。

索引如下

索引

描述

0

输入数据张量。

1

要添加的轴。必须可解析为常量 Int32 或 Int64 1D 形状张量。
参数:

  • index – 输入张量的索引。

  • tensor – 输入张量。

ICumulativeLayer

class tensorrt.ICumulativeLayer

INetworkDefinition 中的累积层。

此层表示跨张量的累积操作。

它计算张量轴上的连续缩减。输出始终具有与输入相同的形状。

如果缩减操作是求和，则这也称为前缀和或累积和。

该操作具有正向与反向变体，以及包含与排除变体。

例如，假设输入是长度为 n 的向量 x，输出是向量 y。则 y[j] = sum(x[…])，其中 … 表示来自此列表的索引序列

  • 包含 + 正向：0..j

  • 包含 + 反向：j..n-1

  • 排除 + 正向：0..j-1

  • 排除 + 反向：j+1..n-1

对于多维张量，累积应用于指定的轴。例如，给定一个 2D 输入，沿轴 0 的正向包含累积会在每列中生成累积和。

变量:

  • op – CumulativeOperation 层的累积操作。

  • exclusive – bool 指定是排除累积还是包含累积。

  • reverse – bool 指定是否应向后应用累积操作。