wehub-resource-sync
718 lines
24 KiB
Python
718 lines
24 KiB
Python
# 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)
|