TensorRT自定义算子¶
TensorRT自定义算子
MMCVRoIAlign
描述
模型参数
输入
输出
类型约束
ScatterND
描述
模型参数
输入
输出
类型约束
NonMaxSuppression
描述
模型参数
输入
输出
类型约束
MMCVDeformConv2d
描述
模型参数
输入
输出
类型约束
grid_sampler
描述
模型参数
输入
输出
类型约束
cummax
描述
模型参数
输入
输出
类型约束
cummin
描述
模型参数
输入
输出
类型约束
MMCVInstanceNormalization
描述
模型参数
输入
输出
类型约束
MMCVModulatedDeformConv2d
描述
模型参数
输入
输出
类型约束
MMCVRoIAlign¶
描述¶
在特征图上计算RoIAlign,在多数双阶段目标检测模型的bbox_head中使用
模型参数¶
类型 | 参数名 | 描述 |
---|---|---|
int |
output_height |
roi特征的输出高度 |
int |
output_width |
roi特征的输出宽度 |
float |
spatial_scale |
输入检测框的缩放系数 |
int |
sampling_ratio |
输出的采样率。0 表示使用密集采样 |
str |
mode |
池化方式。 avg 或max |
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 ,则极值(-1 和1 )会被当做输入边缘像素的中心点。如果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
),其中values
为input
第dim
维的累计最大值,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
),其中values
为input
第dim
维的累计最小值,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)