Shortcuts

TensorRT自定义算子

MMCVRoIAlign

描述

在特征图上计算RoIAlign,在多数双阶段目标检测模型的bbox_head中使用

模型参数

类型 参数名 描述
int output_height roi特征的输出高度
int output_width roi特征的输出宽度
float spatial_scale 输入检测框的缩放系数
int sampling_ratio 输出的采样率。0表示使用密集采样
str mode 池化方式。 avgmax
int aligned 如果aligned=1,则像素会进行-0.5的偏移以达到更好的对齐

输入

inputs[0]: T
输入特征图;形状为(N, C, H, W)的四维张量,其中N为batch大小,C为输入通道数,H和W为输入特征图的高和宽。
inputs[1]: T
需要进行池化的感兴趣区域;形状为(num_rois, 5)的二维张量,内容为[[batch_index, x1, y1, x2, y2], ...]。rois的坐标为输入特征图的坐标系。

输出

outputs[0]: T
池化的输出;形状为(num_rois, C, output_height, output_width)的四维张量。每个输出特征feat[i]都与输入感兴趣区域rois[i]一一对应。
#### 类型约束
  • T:tensor(float32, Linear)

ScatterND

描述

ScatterND接收三个输入,分别为秩为r >= 1的data,秩为q >= 1的indices以及秩为 q + r - indices.shape[-1] -1 的update。输出的计算方式为:首先创建一个data的拷贝,然后根据indces的值使用update对拷贝的data进行更新。注意indices中不应该存在相同的条目,也就是说对同一个位置进行一次以上的更新是不允许的。

输出的计算方式可以参考如下代码:

  output = np.copy(data)
  update_indices = indices.shape[:-1]
  for idx in np.ndindex(update_indices):
      output[indices[idx]] = updates[idx]

模型参数

输入

inputs[0]: T
秩为r >= 1的输入`data`
inputs[1]: tensor(int32, Linear)
秩为q >= 1的输入`update`
inputs[2]: T
秩为 q + r - indices.shape[-1] -1 的输入`update`

输出

outputs[0]: T
秩为r >= 1的输出张量

类型约束

  • T:tensor(float32, Linear), tensor(int32, Linear)

NonMaxSuppression

描述

根据IoU阈值对候选框进行非极大值抑制。

模型参数

类型 参数名 描述
int center_point_box 0 - 候选框的格式为[y1, x1, y2, x2], 1-候选框的格式为[x_center, y_center, width, height]
int max_output_boxes_per_class 每一类最大的输出检测框个数。默认为0,输出检测框个数等于输入候选框数
float iou_threshold 用来判断候选框重合度的阈值,取值范围[0, 1]。默认值为0
float score_threshold 用来判断候选框是否合法的阈值
int offset 检测框长宽计算方式为(x2 - x1 + offset),可选值0或1

输入

inputs[0]: T
输入候选框。形状为(num_batches, spatial_dimension, 4)的三维张量
inputs[1]: T
输入得分。形状为(num_batches, num_classes, spatial_dimension)的三维张量

输出

outputs[0]: tensor(int32, Linear)
被选中的候选框索引。形状为(num_selected_indices, 3)的二维张量。每一行内容为[batch_index, class_index, box_index]。
其中 num_selected_indices=num_batches* num_classes* min(max_output_boxes_per_class, spatial_dimension)。
所有未被选中的候选框索引都会被填充为-1

类型约束

  • T:tensor(float32, Linear)

MMCVDeformConv2d

描述

在输入特征上计算Deformable Convolution,请阅读Deformable Convolutional Network了解更多细节。

模型参数

类型 参数名 描述
list of ints stride 卷积的步长 (sH, sW)
list of ints padding 输入特征填充大小 (padH, padW)
list of ints dilation 卷积核各元素间隔 (dH, dW)
int deformable_group 可变偏移量的分组
int group 卷积分组数,input_channel会根据这个值被分为数个分组进行计算
int im2col_step 可变卷积使用im2col计算卷积。输入与偏移量会以im2col_step为步长分块计算,减少临时空间的使用量。

输入

inputs[0]: T
输入特征;形状为(N, C, inH, inW)的四维张量,其中N为batch大小,C为输入通道数,inH和inW为输入特征图的高和宽
inputs[1]: T
输入偏移量;形状为(N, deformable_group* 2* kH* kW, outH, outW)的四维张量,kH和kW为输入特征图的高和宽,outH和outW为输入特征图的高和宽
inputs[2]: T
输入权重;形状为(output_channel, input_channel, kH, kW)的四维张量

输出

outputs[0]: T
输出特征;形状为(N, output_channel, outH, outW)的四维张量

类型约束

  • T:tensor(float32, Linear)

grid_sampler

描述

根据grid的像素位置对input进行网格采样。

模型参数

类型 参数名 描述
int interpolation_mode 计算输出使用的插值模式。(0: bilinear , 1: nearest)
int padding_mode 边缘填充模式。(0: zeros, 1: border, 2: reflection)
int align_corners 如果align_corners=1,则极值(-11)会被当做输入边缘像素的中心点。如果align_corners=0,则它们会被看做是边缘像素的边缘点,减小分辨率对采样的影响

输入

inputs[0]: T
输入特征;形状为(N, C, inH, inW)的四维张量,其中N为batch大小,C为输入通道数,inH和inW为输入特征图的高和宽
inputs[1]: T
输入网格;形状为(N, outH, outW, 2)的四维张量,outH和outW为输出的高和宽

输出

outputs[0]: T
输出特征;形状为(N, C, outH, outW)的四维张量

类型约束

  • T:tensor(float32, Linear)

cummax

描述

返回一个元组(values, indices),其中valuesinputdim维的累计最大值,indices为第dim维最大值位置。请阅读torch.cummax了解更多细节。

模型参数

类型 参数名 描述
int dim 进行累计计算的维度

输入

inputs[0]: T
输入张量;可以使任意形状

输出

outputs[0]: T
`input`第`dim`维的累计最大值,形状与`input`相同。类型和`input`一致
outputs[1]: (int32, Linear)
第`dim`维最大值位置,形状与`input`相同

类型约束

  • T:tensor(float32, Linear)

cummin

描述

返回一个元组(values, indices),其中valuesinputdim维的累计最小值,indices为第dim维最小值位置。请阅读torch.cummin了解更多细节。

模型参数

类型 参数名 描述
int dim 进行累计计算的维度

输入

inputs[0]: T
输入张量;可以使任意形状

输出

outputs[0]: T
`input`第`dim`维的累计最小值,形状与`input`相同。类型和`input`一致
outputs[1]: (int32, Linear)
第`dim`维最小值位置,形状与`input`相同

类型约束

  • T:tensor(float32, Linear)

MMCVInstanceNormalization

描述

对特征计算instance normalization,请阅读Instance Normalization: The Missing Ingredient for Fast Stylization了解更多详细信息。

模型参数

类型 参数名 描述
float epsilon 用来避免除0错误。默认为1e-05

输入

inputs[0]: T
输入特征。形状为(N, C, H, W)的四维张量,其中N为batch大小,C为输入通道数,H和W为输入特征图的高和宽
inputs[1]: T
输入缩放系数。形状为(C,)的一维张量
inputs[2]: T
输入偏移量。形状为(C,)的一维张量

输出

outputs[0]: T
输出特征。形状为(N, C, H, W)的四维张量

类型约束

  • T:tensor(float32, Linear)

MMCVModulatedDeformConv2d

描述

在输入特征上计算Modulated Deformable Convolution,请阅读Deformable ConvNets v2: More Deformable, Better Results了解更多细节。

模型参数

类型 参数名 描述
list of ints stride 卷积的步长 (sH, sW)
list of ints padding 输入特征填充大小 (padH, padW)
list of ints dilation 卷积核各元素间隔 (dH, dW)
int deformable_groups 可变偏移量的分组,通常置位1即可
int groups 卷积分组数,input_channel会根据这个值被分为数个分组进行计算

输入

inputs[0]: T
输入特征;形状为(N, C, inH, inW)的四维张量,其中N为batch大小,C为输入通道数,inH和inW为输入特征图的高和宽
inputs[1]: T
输入偏移量;形状为(N, deformable_group* 2* kH* kW, outH, outW)的四维张量,kH和kW为输入特征图的高和宽,outH和outW为输入特征图的高和宽
inputs[2]: T
输入掩码;形状为(N, deformable_group* kH* kW, outH, outW)的四维张量
inputs[3]: T
输入权重;形状为(output_channel, input_channel, kH, kW)的四维张量
inputs[4]: T, optional
输入偏移量;形状为(output_channel)的一维张量

输出

outputs[0]: T
输出特征;形状为(N, output_channel, outH, outW)的四维张量

类型约束

  • T:tensor(float32, Linear)

Read the Docs v: v1.3.15
Versions
latest
stable
v1.3.15
v1.3.14
v1.3.13
Downloads
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.