chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,717 @@
|
||||
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
import tensorrt as trt
|
||||
|
||||
from paddle.base.log_helper import get_logger
|
||||
from paddle.tensorrt.converter_utils import (
|
||||
add_1D_constant_layer,
|
||||
fill_constant_layer,
|
||||
get_input_constant_value,
|
||||
get_shape_tensor_element,
|
||||
get_trt_plugin,
|
||||
set_layer_name,
|
||||
trt_concat,
|
||||
trt_div,
|
||||
trt_gather,
|
||||
trt_prod,
|
||||
trt_shape,
|
||||
trt_sub,
|
||||
trt_sum,
|
||||
trt_unsqueeze,
|
||||
)
|
||||
from paddle.tensorrt.register import converter_registry
|
||||
from paddle.tensorrt.util import RefitManager
|
||||
|
||||
_logger = get_logger(
|
||||
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
||||
)
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.multiclass_nms3")
|
||||
def multiclass_nms3_converter(network, paddle_op, inputs):
|
||||
bboxes = inputs[0]
|
||||
scores = inputs[1]
|
||||
background_label = paddle_op.attrs().get("background_label")
|
||||
score_threshold = paddle_op.attrs().get("score_threshold")
|
||||
nms_top_k = paddle_op.attrs().get("nms_top_k")
|
||||
nms_threshold = paddle_op.attrs().get("nms_threshold")
|
||||
keep_top_k = paddle_op.attrs().get("keep_top_k")
|
||||
normalized = paddle_op.attrs().get("normalized")
|
||||
num_classes = scores.shape[1]
|
||||
|
||||
bboxes_dims = bboxes.shape
|
||||
bboxes_expand_dims = [bboxes_dims[0], bboxes_dims[1], 1, bboxes_dims[2]]
|
||||
bboxes_expand_layer = network.add_shuffle(bboxes)
|
||||
bboxes_expand_layer.reshape_dims = trt.Dims(bboxes_expand_dims)
|
||||
set_layer_name(bboxes_expand_layer, paddle_op)
|
||||
|
||||
scores_transpose_layer = network.add_shuffle(scores)
|
||||
scores_transpose_layer.first_transpose = (0, 2, 1)
|
||||
set_layer_name(scores_transpose_layer, paddle_op)
|
||||
|
||||
# create multiclass num3 plugin
|
||||
batch_nms_inputs = [
|
||||
bboxes_expand_layer.get_output(0),
|
||||
scores_transpose_layer.get_output(0),
|
||||
]
|
||||
plugin_fields = [
|
||||
trt.PluginField(
|
||||
"shareLocation",
|
||||
np.array([1], dtype=np.int32),
|
||||
trt.PluginFieldType.INT32,
|
||||
),
|
||||
trt.PluginField(
|
||||
"backgroundLabelId",
|
||||
np.array(background_label, dtype=np.int32),
|
||||
trt.PluginFieldType.INT32,
|
||||
),
|
||||
trt.PluginField(
|
||||
"numClasses",
|
||||
np.array(num_classes, dtype=np.int32),
|
||||
trt.PluginFieldType.INT32,
|
||||
),
|
||||
trt.PluginField(
|
||||
"topK",
|
||||
np.array(nms_top_k, dtype=np.int32),
|
||||
trt.PluginFieldType.INT32,
|
||||
),
|
||||
trt.PluginField(
|
||||
"keepTopK",
|
||||
np.array(keep_top_k, dtype=np.int32),
|
||||
trt.PluginFieldType.INT32,
|
||||
),
|
||||
trt.PluginField(
|
||||
"scoreThreshold",
|
||||
np.array(score_threshold, dtype=np.float32),
|
||||
trt.PluginFieldType.FLOAT32,
|
||||
),
|
||||
trt.PluginField(
|
||||
"iouThreshold",
|
||||
np.array(nms_threshold, dtype=np.float32),
|
||||
trt.PluginFieldType.FLOAT32,
|
||||
),
|
||||
trt.PluginField(
|
||||
"isNormalized",
|
||||
np.array(normalized, dtype=np.int32),
|
||||
trt.PluginFieldType.INT32,
|
||||
),
|
||||
trt.PluginField(
|
||||
"clipBoxes",
|
||||
np.array([0], dtype=np.int32),
|
||||
trt.PluginFieldType.INT32,
|
||||
),
|
||||
]
|
||||
plugin_field_collection = trt.PluginFieldCollection(plugin_fields)
|
||||
plugin_name = "BatchedNMSDynamic_TRT"
|
||||
plugin_version = "1"
|
||||
plugin = get_trt_plugin(
|
||||
plugin_name, plugin_field_collection, plugin_version
|
||||
)
|
||||
batch_nms_layer = network.add_plugin_v2(batch_nms_inputs, plugin)
|
||||
set_layer_name(batch_nms_layer, paddle_op)
|
||||
|
||||
# dynamic shape: [bs, keep_topk, 4], [bs, keep_topk], [bs, keep_topk]
|
||||
nmsed_boxes = batch_nms_layer.get_output(1)
|
||||
nmsed_scores = batch_nms_layer.get_output(2)
|
||||
nmsed_classes = batch_nms_layer.get_output(3)
|
||||
nmsed_scores_transpose_layer = network.add_shuffle(nmsed_scores)
|
||||
set_layer_name(nmsed_scores_transpose_layer, paddle_op)
|
||||
nmsed_classes_reshape_layer = network.add_shuffle(nmsed_classes)
|
||||
set_layer_name(nmsed_classes_reshape_layer, paddle_op)
|
||||
nmsed_scores_transpose_layer.reshape_dims = trt.Dims(
|
||||
[bboxes_dims[0], keep_top_k, 1]
|
||||
)
|
||||
nmsed_classes_reshape_layer.reshape_dims = trt.Dims(
|
||||
[bboxes_dims[0], keep_top_k, 1]
|
||||
)
|
||||
|
||||
concat_inputs = [
|
||||
nmsed_classes_reshape_layer.get_output(0),
|
||||
nmsed_scores_transpose_layer.get_output(0),
|
||||
nmsed_boxes,
|
||||
]
|
||||
nms_concat_layer = network.add_concatenation(inputs=concat_inputs)
|
||||
nms_concat_layer.axis = 2
|
||||
set_layer_name(nms_concat_layer, paddle_op)
|
||||
nms_concat_output = nms_concat_layer.get_output(0)
|
||||
nms_shuffle_layer = network.add_shuffle(nms_concat_output)
|
||||
nms_shuffle_layer.reshape_dims = trt.Dims(
|
||||
[bboxes_dims[0], nms_concat_output.shape[-1]]
|
||||
)
|
||||
set_layer_name(nms_shuffle_layer, paddle_op)
|
||||
|
||||
# add fake index as output to be consistent with the outputs of multiclass_nms3
|
||||
shape_weight = trt.Weights(np.array([0], dtype=np.int32))
|
||||
constant_layer = network.add_constant([1, 1], shape_weight)
|
||||
set_layer_name(constant_layer, paddle_op)
|
||||
|
||||
return (
|
||||
nms_shuffle_layer.get_output(0),
|
||||
constant_layer.get_output(0),
|
||||
batch_nms_layer.get_output(0),
|
||||
)
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.set_value")
|
||||
@converter_registry.register("pd_op.set_value_")
|
||||
@converter_registry.register("pd_op.set_value_with_tensor")
|
||||
@converter_registry.register("pd_op.set_value_with_tensor_")
|
||||
def set_value_converter(network, paddle_op, inputs):
|
||||
x = inputs[0]
|
||||
if (
|
||||
paddle_op.name() == "pd_op.set_value"
|
||||
or paddle_op.name() == "pd_op.set_value_"
|
||||
):
|
||||
starts = get_input_constant_value(paddle_op, inputs, 1)[0]
|
||||
ends = get_input_constant_value(paddle_op, inputs, 2)[0]
|
||||
steps = get_input_constant_value(paddle_op, inputs, 3)[0]
|
||||
else:
|
||||
starts = get_input_constant_value(paddle_op, inputs, 2)[0]
|
||||
ends = get_input_constant_value(paddle_op, inputs, 3)[0]
|
||||
steps = get_input_constant_value(paddle_op, inputs, 4)[0]
|
||||
axes = paddle_op.attrs()["axes"][0]
|
||||
|
||||
input_dims = x.shape
|
||||
|
||||
# check params and refill
|
||||
if axes < 0:
|
||||
axes += len(input_dims)
|
||||
|
||||
if ends < 0:
|
||||
ends += input_dims[axes]
|
||||
|
||||
if ends >= input_dims[axes]:
|
||||
ends = input_dims[axes]
|
||||
|
||||
if (
|
||||
paddle_op.name() == "pd_op.set_value_with_tensor"
|
||||
or paddle_op.name() == "pd_op.set_value_with_tensor_"
|
||||
):
|
||||
updates = inputs[1]
|
||||
else:
|
||||
value = paddle_op.attrs().get("values")
|
||||
input_shape_tensor = trt_shape(
|
||||
network, x, name=[paddle_op.name(), 'input_shape_tensor']
|
||||
)
|
||||
vec_tensor = []
|
||||
for i in range(len(input_dims)):
|
||||
vec_tensor.append(
|
||||
get_shape_tensor_element(
|
||||
network,
|
||||
input_shape_tensor,
|
||||
i,
|
||||
name=[paddle_op.name(), f'vec_tensor_{i}'],
|
||||
)
|
||||
)
|
||||
|
||||
axes_vec = [(ends - 1 - starts) / steps + 1]
|
||||
vec_tensor[axes] = add_1D_constant_layer(
|
||||
network, axes_vec, name=[paddle_op.name(), f'vec_tensor_{axes}']
|
||||
)
|
||||
output_shape_tensor = trt_concat(
|
||||
network,
|
||||
vec_tensor,
|
||||
0,
|
||||
name=[paddle_op.name(), 'output_shape_tensor'],
|
||||
)
|
||||
updates = fill_constant_layer(
|
||||
network,
|
||||
output_shape_tensor,
|
||||
len(x.shape),
|
||||
value,
|
||||
x.dtype,
|
||||
name=[paddle_op.name(), 'updates'],
|
||||
)
|
||||
|
||||
_logger.info(f"Set_value_op: input's dimension is {input_dims}")
|
||||
|
||||
decrease_axes = paddle_op.attrs()["decrease_axes"]
|
||||
if len(decrease_axes) > 0 and len(updates.shape) != len(x.shape):
|
||||
updates = trt_unsqueeze(
|
||||
network,
|
||||
updates,
|
||||
decrease_axes,
|
||||
name=[paddle_op.name(), 'decrease_axes'],
|
||||
)
|
||||
|
||||
value_rank = len(updates.shape)
|
||||
input_rank = len(x.shape)
|
||||
|
||||
assert value_rank == input_rank, (
|
||||
"value's rank is not equal to input's rank, "
|
||||
'you should modify trt_config(a TensorRTConfig object) and set trt_config.disable_ops = ["{op_name}"] to forbid this op '
|
||||
)
|
||||
_logger.info(f"Set_value_op: updates tensor's simension is {updates.shape}")
|
||||
|
||||
# calculate dims
|
||||
update_dims = updates.shape
|
||||
assert update_dims[axes] > 0, (
|
||||
"the update value shape[{axes}] must be greater than 0, but received {update_dims[axes]}"
|
||||
)
|
||||
assert input_dims[axes] > 0, (
|
||||
"the input shape[{axes}] must be greater than 0, but received {input_dims[axes]}"
|
||||
)
|
||||
input_dims_rank = len(input_dims)
|
||||
assert axes <= input_dims_rank, (
|
||||
"The axes {axes} is larger than total axes {input_dims_rank}"
|
||||
)
|
||||
assert starts <= input_dims[axes], (
|
||||
"The start {starts} of dim {axes} is larger than origin shape {input_dims[axes]}"
|
||||
)
|
||||
|
||||
target_update_dim = (ends - 1 - starts) / steps + 1
|
||||
assert update_dims[axes] == target_update_dim, (
|
||||
"the {axes}th axis of update dim error, should be {target_update_dim}, but we got {update_dims[axes]}"
|
||||
)
|
||||
|
||||
shape_0 = [1] * len(update_dims)
|
||||
shape_weight = trt.Weights(np.array([0], dtype=np.float32))
|
||||
zero_tensor = network.add_constant(shape_0, shape_weight)
|
||||
set_layer_name(zero_tensor, paddle_op)
|
||||
zero_tensor = zero_tensor.get_output(0)
|
||||
|
||||
indice_tensor = trt_prod(
|
||||
network, zero_tensor, updates, name=[paddle_op.name(), 'indice_tensor']
|
||||
)
|
||||
cast_layer = network.add_identity(indice_tensor)
|
||||
set_layer_name(cast_layer, paddle_op)
|
||||
cast_layer.set_output_type(0, trt.int32)
|
||||
indice_tensor = cast_layer.get_output(0)
|
||||
|
||||
shape_1 = [1] * len(update_dims)
|
||||
shape_1[axes] = update_dims[axes]
|
||||
tmp_1 = []
|
||||
for i in range(starts, ends, steps):
|
||||
tmp_1.append(i)
|
||||
shape_weight = trt.Weights(np.array(tmp_1, dtype=np.int32))
|
||||
one_tensor = network.add_constant(shape_1, shape_weight)
|
||||
set_layer_name(one_tensor, paddle_op)
|
||||
one_tensor = one_tensor.get_output(0)
|
||||
|
||||
indice_tensor = trt_sum(
|
||||
network,
|
||||
indice_tensor,
|
||||
one_tensor,
|
||||
name=[paddle_op.name(), 'indice_tensor'],
|
||||
)
|
||||
layer = network.add_scatter(
|
||||
x, indice_tensor, updates, trt.ScatterMode.ELEMENT
|
||||
)
|
||||
set_layer_name(layer, paddle_op)
|
||||
layer.axis = axes
|
||||
return layer.get_output(0)
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.share_data")
|
||||
@converter_registry.register("pd_op.share_data_")
|
||||
def share_data_converter(network, paddle_op, inputs):
|
||||
x = inputs[0]
|
||||
identity_layer = network.add_identity(x)
|
||||
set_layer_name(identity_layer, paddle_op)
|
||||
return identity_layer.get_output(0)
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.temporal_shift")
|
||||
def temporal_shift_converter(network, paddle_op, inputs):
|
||||
input_tensor = inputs[0]
|
||||
# Add a small bias to shift_ratio to mitigate floating point precision errors
|
||||
shift_ratio = paddle_op.attrs()["shift_ratio"] + 1e-7
|
||||
T = paddle_op.attrs()["seg_num"]
|
||||
data_format = paddle_op.attrs().get("data_format", "NCHW")
|
||||
|
||||
if data_format == "NHWC":
|
||||
# Transpose input to [N, C, H, W]
|
||||
transpose_layer = network.add_shuffle(input_tensor)
|
||||
transpose_layer.first_transpose = trt.Permutation([0, 3, 1, 2])
|
||||
set_layer_name(transpose_layer, paddle_op)
|
||||
input_tensor = transpose_layer.get_output(0)
|
||||
|
||||
input_dims = input_tensor.shape
|
||||
C, H, W = input_dims[1], input_dims[2], input_dims[3]
|
||||
|
||||
# Reshape input to [N, T, C, H, W]
|
||||
reshape_layer = network.add_shuffle(input_tensor)
|
||||
reshape_layer.reshape_dims = trt.Dims([-1, T, C, H, W])
|
||||
set_layer_name(reshape_layer, paddle_op)
|
||||
input_tensor = reshape_layer.get_output(0)
|
||||
|
||||
# Pad input to [N, T + 2, C, H, W]
|
||||
pre_pad = add_1D_constant_layer(
|
||||
network, [0, 1, 0, 0, 0], name=[paddle_op.name(), 'pre_pad']
|
||||
)
|
||||
post_pad = add_1D_constant_layer(
|
||||
network, [0, 1, 0, 0, 0], name=[paddle_op.name(), 'post_pad']
|
||||
)
|
||||
dims = 5
|
||||
zeros = add_1D_constant_layer(
|
||||
network, [0] * dims, name=[paddle_op.name(), 'zeros']
|
||||
)
|
||||
start = trt_sub(network, zeros, pre_pad, name=[paddle_op.name(), 'start'])
|
||||
total_padding = trt_sum(
|
||||
network, pre_pad, post_pad, name=[paddle_op.name(), 'total_padding']
|
||||
)
|
||||
input_shape = trt_shape(
|
||||
network, input_tensor, name=[paddle_op.name(), 'input_shape']
|
||||
)
|
||||
size = trt_sum(
|
||||
network, input_shape, total_padding, name=[paddle_op.name(), 'size']
|
||||
)
|
||||
stride = [1] * dims
|
||||
dummy = stride
|
||||
|
||||
slice_layer = network.add_slice(input_tensor, dummy, dummy, stride)
|
||||
slice_layer.set_input(1, start)
|
||||
slice_layer.set_input(2, size)
|
||||
set_layer_name(slice_layer, paddle_op)
|
||||
|
||||
trt_version = trt.__version__.split('.')
|
||||
if int(trt_version[0]) > 8 or (
|
||||
int(trt_version[0]) == 8 and int(trt_version[1]) >= 5
|
||||
):
|
||||
slice_layer.mode = trt.SampleMode.FILL
|
||||
else:
|
||||
slice_layer.mode = trt.SliceMode.FILL
|
||||
|
||||
slice_c = int(C * shift_ratio)
|
||||
slice_c2 = int(C * shift_ratio * 2)
|
||||
|
||||
slice_start1 = zeros
|
||||
slice_start2 = add_1D_constant_layer(
|
||||
network, [0, 2, slice_c, 0, 0], name=[paddle_op.name(), 'slice_start2']
|
||||
)
|
||||
slice_start3 = add_1D_constant_layer(
|
||||
network, [0, 1, slice_c2, 0, 0], name=[paddle_op.name(), 'slice_start3']
|
||||
)
|
||||
|
||||
slice_size_base = trt_shape(
|
||||
network, input_tensor, name=[paddle_op.name(), 'slice_size_base']
|
||||
)
|
||||
sub_size1 = add_1D_constant_layer(
|
||||
network, [0, 0, C - slice_c, 0, 0], name=[paddle_op.name(), 'sub_size1']
|
||||
)
|
||||
sub_size2 = add_1D_constant_layer(
|
||||
network,
|
||||
[0, 0, C + slice_c - slice_c2, 0, 0],
|
||||
name=[paddle_op.name(), 'sub_size2'],
|
||||
)
|
||||
sub_size3 = add_1D_constant_layer(
|
||||
network, [0, 0, slice_c2, 0, 0], name=[paddle_op.name(), 'sub_size3']
|
||||
)
|
||||
|
||||
slice_size1 = trt_sub(
|
||||
network,
|
||||
slice_size_base,
|
||||
sub_size1,
|
||||
name=[paddle_op.name(), 'slice_size1'],
|
||||
)
|
||||
slice_size2 = trt_sub(
|
||||
network,
|
||||
slice_size_base,
|
||||
sub_size2,
|
||||
name=[paddle_op.name(), 'slice_size2'],
|
||||
)
|
||||
slice_size3 = trt_sub(
|
||||
network,
|
||||
slice_size_base,
|
||||
sub_size3,
|
||||
name=[paddle_op.name(), 'slice_size3'],
|
||||
)
|
||||
|
||||
slice1_layer = network.add_slice(
|
||||
slice_layer.get_output(0), start=dummy, shape=dummy, stride=stride
|
||||
)
|
||||
slice1_layer.set_input(1, slice_start1)
|
||||
slice1_layer.set_input(2, slice_size1)
|
||||
set_layer_name(slice1_layer, paddle_op)
|
||||
slice2_layer = network.add_slice(
|
||||
slice_layer.get_output(0), start=dummy, shape=dummy, stride=stride
|
||||
)
|
||||
slice2_layer.set_input(1, slice_start2)
|
||||
slice2_layer.set_input(2, slice_size2)
|
||||
set_layer_name(slice2_layer, paddle_op)
|
||||
slice3_layer = network.add_slice(
|
||||
slice_layer.get_output(0), start=dummy, shape=dummy, stride=stride
|
||||
)
|
||||
slice3_layer.set_input(1, slice_start3)
|
||||
slice3_layer.set_input(2, slice_size3)
|
||||
set_layer_name(slice3_layer, paddle_op)
|
||||
|
||||
concat_inputs = [slice2_layer.get_output(0), slice3_layer.get_output(0)]
|
||||
if slice_c == 0:
|
||||
concat_layer = network.add_concatenation(concat_inputs)
|
||||
concat_layer.axis = 2
|
||||
set_layer_name(concat_layer, paddle_op)
|
||||
else:
|
||||
concat_inputs = [
|
||||
slice1_layer.get_output(0),
|
||||
slice2_layer.get_output(0),
|
||||
slice3_layer.get_output(0),
|
||||
]
|
||||
concat_layer = network.add_concatenation(concat_inputs)
|
||||
concat_layer.axis = 2
|
||||
set_layer_name(concat_layer, paddle_op)
|
||||
|
||||
# Reshape output to [N*T,C,H,W]
|
||||
reshape_layer3 = network.add_shuffle(concat_layer.get_output(0))
|
||||
reshape_layer3.reshape_dims = trt.Dims([-1, C, H, W])
|
||||
set_layer_name(reshape_layer3, paddle_op)
|
||||
|
||||
if data_format == "NHWC":
|
||||
transpose_layer2 = network.add_shuffle(reshape_layer3.get_output(0))
|
||||
transpose_layer2.first_transpose = trt.Permutation([0, 2, 3, 1])
|
||||
set_layer_name(transpose_layer2, paddle_op)
|
||||
output_tensor = transpose_layer2.get_output(0)
|
||||
else:
|
||||
output_tensor = reshape_layer3.get_output(0)
|
||||
|
||||
return output_tensor
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.anchor_generator")
|
||||
def anchor_generator_converter(network, paddle_op, inputs):
|
||||
inputs = inputs[0]
|
||||
input_dims = inputs.shape
|
||||
anchor_sizes = paddle_op.attrs().get("anchor_sizes")
|
||||
aspect_ratios = paddle_op.attrs().get("aspect_ratios")
|
||||
stride = paddle_op.attrs().get("stride")
|
||||
variances = paddle_op.attrs().get("variances")
|
||||
offset = paddle_op.attrs().get("offset")
|
||||
num_anchors = len(aspect_ratios) * len(anchor_sizes)
|
||||
|
||||
height = input_dims[1]
|
||||
width = input_dims[2]
|
||||
|
||||
plugin_fields = [
|
||||
trt.PluginField(
|
||||
"anchor_sizes",
|
||||
np.array(anchor_sizes, dtype=np.float32),
|
||||
trt.PluginFieldType.FLOAT32,
|
||||
),
|
||||
trt.PluginField(
|
||||
"aspect_ratios",
|
||||
np.array(aspect_ratios, dtype=np.float32),
|
||||
trt.PluginFieldType.FLOAT32,
|
||||
),
|
||||
trt.PluginField(
|
||||
"stride",
|
||||
np.array(stride, dtype=np.float32),
|
||||
trt.PluginFieldType.FLOAT32,
|
||||
),
|
||||
trt.PluginField(
|
||||
"variances",
|
||||
np.array(variances, dtype=np.float32),
|
||||
trt.PluginFieldType.FLOAT32,
|
||||
),
|
||||
trt.PluginField(
|
||||
"offset",
|
||||
np.array(offset, dtype=np.float32),
|
||||
trt.PluginFieldType.FLOAT32,
|
||||
),
|
||||
trt.PluginField(
|
||||
"num_anchors",
|
||||
np.array(num_anchors, dtype=np.int32),
|
||||
trt.PluginFieldType.INT32,
|
||||
),
|
||||
]
|
||||
plugin_field_collection = trt.PluginFieldCollection(plugin_fields)
|
||||
plugin_name = "pir_anchor_generator_plugin_dynamic"
|
||||
plugin_version = "1"
|
||||
plugin = get_trt_plugin(
|
||||
plugin_name, plugin_field_collection, plugin_version
|
||||
)
|
||||
anchor_generator_layer = network.add_plugin_v2([inputs], plugin)
|
||||
set_layer_name(anchor_generator_layer, paddle_op)
|
||||
out0 = anchor_generator_layer.get_output(0)
|
||||
out1 = anchor_generator_layer.get_output(1)
|
||||
return (out0, out1)
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.affine_channel")
|
||||
def affine_channel_converter(network, paddle_op, inputs):
|
||||
x, scale, bias = inputs
|
||||
data_layout = paddle_op.attrs().get("data_layout")
|
||||
if isinstance(scale, trt.ITensor):
|
||||
refit_manager = RefitManager()
|
||||
scale_weights = refit_manager.get_trt_weight_tensor(scale.name)
|
||||
bias_weights = refit_manager.get_trt_weight_tensor(bias.name)
|
||||
else:
|
||||
scale_weights = scale
|
||||
bias_weights = bias
|
||||
|
||||
if data_layout == "NCHW":
|
||||
channel_axis = 1
|
||||
x_input = x
|
||||
elif data_layout == "NHWC":
|
||||
# Permute NHWC to NCHW
|
||||
shuffle_layer1 = network.add_shuffle(x)
|
||||
shuffle_layer1.first_transpose = (0, 3, 1, 2)
|
||||
set_layer_name(shuffle_layer1, paddle_op)
|
||||
x_input = shuffle_layer1.get_output(0)
|
||||
channel_axis = 1
|
||||
else:
|
||||
raise ValueError(f"affine_channel: Unsupported layout: {data_layout}")
|
||||
|
||||
if scale_weights.size != bias_weights.size:
|
||||
raise ValueError(
|
||||
f"affine_channel: scale.size({scale_weights.size}) != bias.size({bias_weights.size})"
|
||||
)
|
||||
|
||||
power_array = np.ones((scale_weights.size,), dtype=np.float32)
|
||||
power_weights = trt.Weights(power_array)
|
||||
|
||||
layer = network.add_scale_nd(
|
||||
input=x_input,
|
||||
mode=trt.ScaleMode.CHANNEL,
|
||||
shift=bias_weights,
|
||||
scale=scale_weights,
|
||||
power=power_weights,
|
||||
channel_axis=channel_axis,
|
||||
)
|
||||
set_layer_name(layer, paddle_op)
|
||||
if not layer:
|
||||
raise RuntimeError("affine_channel: add_scale_nd failed.")
|
||||
|
||||
out_tensor = layer.get_output(0)
|
||||
|
||||
if data_layout == "NHWC":
|
||||
shuffle_layer2 = network.add_shuffle(out_tensor)
|
||||
shuffle_layer2.first_transpose = (0, 2, 3, 1)
|
||||
set_layer_name(shuffle_layer2, paddle_op)
|
||||
out_tensor = shuffle_layer2.get_output(0)
|
||||
|
||||
return out_tensor
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.shuffle_channel")
|
||||
def shuffle_channel_converter(network, paddle_op, inputs):
|
||||
input = inputs[0]
|
||||
group = paddle_op.attrs().get("group")
|
||||
input_shape_tensor = trt_shape(
|
||||
network, input, name=[paddle_op.name(), 'input_shape_tensor']
|
||||
)
|
||||
batch_shape_tensor = get_shape_tensor_element(
|
||||
network,
|
||||
input_shape_tensor,
|
||||
0,
|
||||
name=[paddle_op.name(), 'batch_shape_tensor'],
|
||||
)
|
||||
channel_shape_tensor = get_shape_tensor_element(
|
||||
network,
|
||||
input_shape_tensor,
|
||||
1,
|
||||
name=[paddle_op.name(), 'channel_shape_tensor'],
|
||||
)
|
||||
group_tensor = add_1D_constant_layer(
|
||||
network, group, name=[paddle_op.name(), 'group_tensor']
|
||||
)
|
||||
new_channel_shape_tensor = trt_div(
|
||||
network,
|
||||
channel_shape_tensor,
|
||||
group_tensor,
|
||||
name=[paddle_op.name(), 'new_channel_shape_tensor'],
|
||||
)
|
||||
shape_dim2 = [2, 3]
|
||||
shape_dim2_tensor = trt_gather(
|
||||
network,
|
||||
input_shape_tensor,
|
||||
shape_dim2,
|
||||
name=[paddle_op.name(), 'shape_dim2_tensor'],
|
||||
)
|
||||
itensors = [
|
||||
batch_shape_tensor,
|
||||
group_tensor,
|
||||
new_channel_shape_tensor,
|
||||
shape_dim2_tensor,
|
||||
]
|
||||
reshape_tensor = trt_concat(
|
||||
network, itensors, name=[paddle_op.name(), 'reshape_tensor']
|
||||
)
|
||||
layer = network.add_shuffle(input)
|
||||
layer.set_input(1, reshape_tensor)
|
||||
transpose_embed = trt.Permutation([0, 2, 1, 3, 4])
|
||||
layer.second_transpose = transpose_embed
|
||||
set_layer_name(layer, paddle_op)
|
||||
output = layer.get_output(0)
|
||||
output_layer = network.add_shuffle(output)
|
||||
output_layer.set_input(1, input_shape_tensor)
|
||||
set_layer_name(output_layer, paddle_op)
|
||||
return output_layer.get_output(0)
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.full_batch_size_like")
|
||||
def full_batch_size_like_converter(network, paddle_op, inputs):
|
||||
input = inputs[0]
|
||||
input_dim_idx = paddle_op.attrs().get("input_dim_idx")
|
||||
output_dim_idx = paddle_op.attrs().get("output_dim_idx")
|
||||
value = paddle_op.attrs().get("value")
|
||||
shape = paddle_op.attrs().get("shape")
|
||||
value = float(value)
|
||||
|
||||
input_shape_tensor = trt_shape(
|
||||
network, input, name=[paddle_op.name(), 'input_shape_tensor']
|
||||
)
|
||||
batch_tensor = get_shape_tensor_element(
|
||||
network,
|
||||
input_shape_tensor,
|
||||
input_dim_idx,
|
||||
name=[paddle_op.name(), 'batch_tensor'],
|
||||
)
|
||||
|
||||
shape_attr_tensor = add_1D_constant_layer(
|
||||
network, shape, name=[paddle_op.name(), 'shape_attr_tensor']
|
||||
)
|
||||
|
||||
gather_output_shape_indices = [
|
||||
len(shape) if i == output_dim_idx else i for i in range(len(shape))
|
||||
]
|
||||
|
||||
concat_inputs = [shape_attr_tensor, batch_tensor]
|
||||
concat_tensor = trt_concat(
|
||||
network, concat_inputs, name=[paddle_op.name(), 'concat_tensor']
|
||||
)
|
||||
out_shape_tensor = trt_gather(
|
||||
network,
|
||||
concat_tensor,
|
||||
gather_output_shape_indices,
|
||||
name=[paddle_op.name(), 'out_shape_tensor'],
|
||||
)
|
||||
|
||||
layer = network.add_fill(shape=(), op=trt.FillOperation.LINSPACE)
|
||||
|
||||
value_tensor = add_1D_constant_layer(
|
||||
network,
|
||||
[value],
|
||||
is_scalar=True,
|
||||
name=[paddle_op.name(), 'value_tensor'],
|
||||
)
|
||||
|
||||
beta_vec = [0.0] * len(shape)
|
||||
beta_tensor = add_1D_constant_layer(
|
||||
network,
|
||||
beta_vec,
|
||||
is_scalar=False,
|
||||
name=[paddle_op.name(), 'beta_tensor'],
|
||||
)
|
||||
|
||||
layer.set_input(0, out_shape_tensor)
|
||||
layer.set_input(1, value_tensor)
|
||||
layer.set_input(2, beta_tensor)
|
||||
|
||||
set_layer_name(layer, paddle_op)
|
||||
|
||||
return layer.get_output(0)
|
||||
Reference in New Issue
Block a user