wehub-resource-sync
1634 lines
52 KiB
Python
1634 lines
52 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 numpy as np
|
|
import tensorrt as trt
|
|
|
|
from paddle.tensorrt.converter_utils import (
|
|
add_1D_constant_layer,
|
|
build_size_tensor,
|
|
build_start_tensor,
|
|
cast_tensor,
|
|
fix_negative_indices,
|
|
generic_plugin_converter,
|
|
get_axes_for_reduce_op,
|
|
get_input_constant_value,
|
|
get_shape_tensor_element,
|
|
has_dynamic_shape,
|
|
resize_to_1d,
|
|
set_layer_name,
|
|
trt_cast,
|
|
trt_concat,
|
|
trt_expand,
|
|
trt_floor_div,
|
|
trt_gather,
|
|
trt_less,
|
|
trt_max,
|
|
trt_min,
|
|
trt_prod,
|
|
trt_reshape,
|
|
trt_shape,
|
|
trt_sub,
|
|
trt_sum,
|
|
)
|
|
from paddle.tensorrt.register import converter_registry
|
|
|
|
from ..util import get_trt_version_list
|
|
|
|
|
|
@converter_registry.register("pd_op.reshape")
|
|
def reshape_converter(network, paddle_op, inputs):
|
|
x = inputs[0]
|
|
is_constant_shape = False
|
|
shape = get_input_constant_value(paddle_op, inputs, 1)
|
|
if shape is not None:
|
|
reshape_dim = shape
|
|
is_constant_shape = True
|
|
elif isinstance(inputs[1], list):
|
|
# shape tensor is a list value
|
|
shape_tensor = trt_concat(
|
|
network, inputs[1], name=[paddle_op.name(), "shape_tensor"]
|
|
)
|
|
else:
|
|
# shape tensor is a value
|
|
shape_tensor = inputs[1]
|
|
|
|
if not is_constant_shape:
|
|
shape_tensor = resize_to_1d(
|
|
network, shape_tensor, name=[paddle_op.name(), "shape_tensor"]
|
|
)
|
|
|
|
layer = network.add_shuffle(x)
|
|
if is_constant_shape:
|
|
layer.reshape_dims = reshape_dim
|
|
else:
|
|
layer.set_input(1, shape_tensor)
|
|
set_layer_name(layer, paddle_op)
|
|
|
|
assert len(layer.get_output(0).shape) >= 0, (
|
|
'When convert reshape op to TRT reshape layer, the rank of trt reshape output dims is less than 0, '
|
|
'you should modify trt_config(a TensorRTConfig object) and set trt_config.disable_ops = ["pd_op.reshape"] to forbid this op.'
|
|
)
|
|
|
|
return layer.get_output(0)
|
|
|
|
|
|
@converter_registry.register("pd_op.gather")
|
|
def gather_converter(network, paddle_op, inputs):
|
|
input_tensor = inputs[0]
|
|
index_tensor = inputs[1]
|
|
axis_value = get_input_constant_value(paddle_op, inputs, 2)[0]
|
|
axis_value = int(axis_value)
|
|
|
|
reshape_layer = network.add_shuffle(index_tensor)
|
|
reshape_layer.reshape_dims = (-1,)
|
|
set_layer_name(reshape_layer, paddle_op)
|
|
gather_layer = network.add_gather(
|
|
input_tensor, reshape_layer.get_output(0), axis_value
|
|
)
|
|
set_layer_name(gather_layer, paddle_op)
|
|
return gather_layer.get_output(0)
|
|
|
|
|
|
@converter_registry.register("pd_op.gather_nd")
|
|
def gather_nd_converter(network, paddle_op, inputs):
|
|
input_tensor, indices_tensor = inputs
|
|
non_zero_layer = network.add_gather_v2(
|
|
input_tensor, indices_tensor, trt.GatherMode.ND
|
|
)
|
|
non_zero_layer.num_elementwise_dims = 0
|
|
set_layer_name(non_zero_layer, paddle_op)
|
|
return non_zero_layer.get_output(0)
|
|
|
|
|
|
@converter_registry.register("pd_op.flatten")
|
|
def flatten_converter(network, paddle_op, inputs):
|
|
input_val = inputs[0]
|
|
input_val_shape = paddle_op.operands()[0].source().shape
|
|
dims = len(input_val_shape)
|
|
|
|
start_axis = paddle_op.attrs().get("start_axis")
|
|
stop_axis = paddle_op.attrs().get("stop_axis")
|
|
|
|
flatten_layer = network.add_shuffle(input_val)
|
|
set_layer_name(flatten_layer, paddle_op)
|
|
|
|
if not has_dynamic_shape(input_val_shape):
|
|
if start_axis < 0:
|
|
start_axis += dims + 1
|
|
if stop_axis < 0:
|
|
stop_axis += dims + 1
|
|
|
|
flatten_dim = 1
|
|
final_shape = []
|
|
|
|
for i, s in enumerate(input_val_shape):
|
|
if i >= start_axis and i <= stop_axis:
|
|
flatten_dim *= s
|
|
elif i == stop_axis + 1:
|
|
final_shape.append(flatten_dim)
|
|
final_shape.append(s)
|
|
else:
|
|
final_shape.append(s)
|
|
|
|
if stop_axis == len(input_val.shape) - 1:
|
|
final_shape.append(flatten_dim)
|
|
|
|
flatten_layer.reshape_dims = tuple(final_shape)
|
|
set_layer_name(flatten_layer, paddle_op)
|
|
else:
|
|
input_shape_layer = network.add_shape(input_val)
|
|
set_layer_name(input_shape_layer, paddle_op)
|
|
|
|
final_shapes = []
|
|
# Shapes before start_axis
|
|
if start_axis > 0:
|
|
prefix_shape_layer = network.add_slice(
|
|
input_shape_layer.get_output(0),
|
|
start=(0,),
|
|
shape=(start_axis,),
|
|
stride=(1,),
|
|
)
|
|
set_layer_name(prefix_shape_layer, paddle_op)
|
|
final_shapes.append(prefix_shape_layer.get_output(0))
|
|
|
|
flatten_shape_layer = network.add_slice(
|
|
input_shape_layer.get_output(0),
|
|
start=(start_axis,),
|
|
shape=(stop_axis - start_axis + 1,),
|
|
stride=(1,),
|
|
)
|
|
set_layer_name(flatten_shape_layer, paddle_op)
|
|
flatten_shape_layer = network.add_reduce(
|
|
flatten_shape_layer.get_output(0),
|
|
trt.ReduceOperation.PROD,
|
|
axes=get_axes_for_reduce_op(0, False),
|
|
keep_dims=True,
|
|
)
|
|
set_layer_name(flatten_shape_layer, paddle_op)
|
|
final_shapes.append(flatten_shape_layer.get_output(0))
|
|
|
|
# Shapes after stop_axis
|
|
if stop_axis < len(input_val_shape) - 1:
|
|
suffix_shape_layer = network.add_slice(
|
|
input_shape_layer.get_output(0),
|
|
start=(stop_axis + 1,),
|
|
shape=(len(input_val_shape) - stop_axis - 1,),
|
|
stride=(1,),
|
|
)
|
|
set_layer_name(suffix_shape_layer, paddle_op)
|
|
final_shapes.append(suffix_shape_layer.get_output(0))
|
|
|
|
final_shape_layer = network.add_concatenation(final_shapes)
|
|
final_shape_layer.axis = 0
|
|
set_layer_name(final_shape_layer, paddle_op)
|
|
flatten_layer.set_input(1, final_shape_layer.get_output(0))
|
|
|
|
return flatten_layer.get_output(0)
|
|
|
|
|
|
# In the converter, pd_op.concat has three inputs, because builtin.combine has two inputs.
|
|
@converter_registry.register("pd_op.concat")
|
|
def concat_converter(network, paddle_op, inputs):
|
|
input_tensors = inputs[0]
|
|
concat_layer = network.add_concatenation(inputs=input_tensors)
|
|
|
|
axis = get_input_constant_value(paddle_op, inputs, 1)[0]
|
|
axis = int(axis)
|
|
if axis < 0:
|
|
axis = len(input_tensors[0].shape) + axis
|
|
concat_layer.axis = axis
|
|
set_layer_name(concat_layer, paddle_op)
|
|
|
|
return concat_layer.get_output(0)
|
|
|
|
|
|
@converter_registry.register("pd_op.unsqueeze")
|
|
@converter_registry.register("pd_op.unsqueeze_")
|
|
def unsqueeze_converter(network, paddle_op, inputs):
|
|
x = inputs[0]
|
|
input_dims = x.shape
|
|
axes = get_input_constant_value(paddle_op, inputs, 1)
|
|
assert len(axes) > 0, (
|
|
f"axes size should be > 0 in when convert unsqueeze op in TensorRT, but received len(axes) = {len(axes)}."
|
|
)
|
|
|
|
should_unsqueeze = [False] * (len(input_dims) + len(axes))
|
|
cur_out_rank = len(input_dims)
|
|
for i in range(len(axes)):
|
|
cur_out_rank += 1
|
|
if axes[i] < 0:
|
|
axes[i] += cur_out_rank
|
|
|
|
# axes[i] is relative to cur_out_rank
|
|
# we make [axes[i], cur_out_rank - 2] shift right
|
|
# and make (axes[i]) to true!
|
|
for j in range(cur_out_rank - 1, axes[i], -1):
|
|
should_unsqueeze[j] = should_unsqueeze[j - 1]
|
|
if axes[i] >= cur_out_rank:
|
|
should_unsqueeze[cur_out_rank - 1] = True
|
|
else:
|
|
should_unsqueeze[axes[i]] = True
|
|
|
|
gather_indices = []
|
|
in_rank_i = 0
|
|
for i in range(len(should_unsqueeze)):
|
|
if should_unsqueeze[i]:
|
|
gather_indices.append(len(input_dims))
|
|
continue
|
|
gather_indices.append(in_rank_i)
|
|
in_rank_i += 1
|
|
|
|
shape_tensor = trt_shape(
|
|
network, x, name=[paddle_op.name(), "shape_tensor"]
|
|
)
|
|
all_one = [1] * len(axes)
|
|
all_one_tensor = add_1D_constant_layer(
|
|
network, all_one, name=[paddle_op.name(), "all_one_tensor"]
|
|
)
|
|
concat_inputs = [shape_tensor, all_one_tensor]
|
|
real_shape_tensor = trt_gather(
|
|
network,
|
|
trt_concat(
|
|
network, concat_inputs, name=[paddle_op.name(), "trt_concat"]
|
|
),
|
|
gather_indices,
|
|
name=[paddle_op.name(), "real_shape_tensor"],
|
|
)
|
|
layer = network.add_shuffle(x)
|
|
layer.set_input(1, real_shape_tensor)
|
|
set_layer_name(layer, paddle_op)
|
|
return layer.get_output(0)
|
|
|
|
|
|
@converter_registry.register("pd_op.squeeze")
|
|
@converter_registry.register("pd_op.squeeze_")
|
|
def squeeze_converter(network, paddle_op, inputs):
|
|
input_val = inputs[0]
|
|
input_shape = input_val.shape
|
|
input_shape_size = len(input_shape)
|
|
|
|
# If input is weights, convert to TensorRT tensor
|
|
if isinstance(input_val, trt.Weights):
|
|
input_val = network.add_constant(input_shape, input_val)
|
|
set_layer_name(input_val, paddle_op)
|
|
input_val = input_val.get_output(0)
|
|
|
|
# Get axis
|
|
axis = get_input_constant_value(paddle_op, inputs, 1)
|
|
if len(axis) == 0:
|
|
for i in range(input_shape_size):
|
|
if input_shape[i] == -1:
|
|
raise RuntimeError(
|
|
"The necessary attributes of the squeeze operator axis is missing"
|
|
)
|
|
elif input_shape[i] == 1:
|
|
axis.append(i)
|
|
else:
|
|
# Verify that each axis to squeeze has size 1
|
|
for a in axis:
|
|
if a < 0:
|
|
a += input_shape_size
|
|
if input_shape[a] != 1:
|
|
raise RuntimeError(
|
|
f"Cannot squeeze dimension {a} with size {input_shape[a]}. Only dimensions with size 1 can be squeezed."
|
|
)
|
|
|
|
axes_size = len(axis)
|
|
if axes_size == 0:
|
|
raise RuntimeError(
|
|
f"axis.size should be >0 in pd_op.squeeze op in TensorRT, but received {axes_size}"
|
|
)
|
|
# Mark which dimensions to squeeze
|
|
should_squeeze = [False] * input_shape_size
|
|
for a in axis:
|
|
should_squeeze[a] = True
|
|
|
|
# Get dimensions to keep
|
|
gather_indices = [
|
|
i for i, squeeze in enumerate(should_squeeze) if not squeeze
|
|
]
|
|
|
|
# Add Shuffle layer
|
|
shape_tensor = trt_shape(
|
|
network, input_val, name=[paddle_op.name(), 'shape_tensor']
|
|
)
|
|
real_shape_tensor = trt_gather(
|
|
network,
|
|
shape_tensor,
|
|
gather_indices,
|
|
name=[paddle_op.name(), 'real_shape_tensor'],
|
|
)
|
|
layer = network.add_shuffle(input_val)
|
|
layer.set_input(1, real_shape_tensor)
|
|
set_layer_name(layer, paddle_op)
|
|
|
|
return layer.get_output(0)
|
|
|
|
|
|
@converter_registry.register("pd_op.expand")
|
|
def expand_converter(network, paddle_op, inputs):
|
|
input = inputs[0]
|
|
input_dims = input.shape
|
|
rank = len(input_dims)
|
|
paddle_shape_tensor = paddle_op.operands()[1].source()
|
|
|
|
shape = get_input_constant_value(paddle_op, inputs, 1)
|
|
if shape is not None:
|
|
shape_tensor = add_1D_constant_layer(
|
|
network, shape, name=[paddle_op.name(), 'shape_tensor']
|
|
)
|
|
shape_rank = len(shape)
|
|
elif paddle_shape_tensor.type().as_vec_type():
|
|
shape_tensors = inputs[1]
|
|
shape_rank = len(shape_tensors)
|
|
shape_tensor = trt_concat(
|
|
network, shape_tensors, name=[paddle_op.name(), 'shape_tensor']
|
|
)
|
|
else:
|
|
shape_tensor = inputs[1]
|
|
shape_rank = shape_tensor.shape[0]
|
|
return trt_expand(
|
|
network,
|
|
input,
|
|
rank,
|
|
shape_tensor,
|
|
shape_rank,
|
|
name=[paddle_op.name(), 'trt_expand'],
|
|
)
|
|
|
|
|
|
@converter_registry.register("pd_op.expand_as")
|
|
def expand_as_converter(network, paddle_op, inputs):
|
|
input = inputs[0]
|
|
input_dims = input.shape
|
|
rank = len(input_dims)
|
|
y = paddle_op.operands()[1].source()
|
|
|
|
if y.initialized():
|
|
y_t = inputs[1]
|
|
shape_tensor = trt_shape(
|
|
network, y_t, name=[paddle_op.name(), 'shape_tensor']
|
|
)
|
|
shape_rank = len(y_t.shape)
|
|
else:
|
|
shape = paddle_op.attrs().get("target_shape")
|
|
shape_tensor = add_1D_constant_layer(
|
|
network, shape, name=[paddle_op.name(), 'shape_tensor']
|
|
)
|
|
shape_rank = len(shape)
|
|
return trt_expand(
|
|
network,
|
|
input,
|
|
rank,
|
|
shape_tensor,
|
|
shape_rank,
|
|
name=[paddle_op.name(), 'trt_expand'],
|
|
)
|
|
|
|
|
|
@converter_registry.register("pd_op.cast")
|
|
@converter_registry.register("pd_op.cast_")
|
|
def cast_converter(network, paddle_op, inputs):
|
|
input_tensor = inputs[0]
|
|
out_dtype = int(paddle_op.attrs().get("dtype"))
|
|
# Reference paddle/phi/common/data_type.h enum DataType
|
|
if out_dtype == 1:
|
|
out_dtype = trt.bool
|
|
elif out_dtype == 7:
|
|
out_dtype = trt.int32
|
|
elif out_dtype == 9:
|
|
out_dtype = trt.int32
|
|
elif out_dtype == 10:
|
|
out_dtype = trt.float32
|
|
elif out_dtype == 11:
|
|
out_dtype = trt.float32
|
|
elif out_dtype == 15:
|
|
out_dtype = trt.float16
|
|
else:
|
|
raise RuntimeError(
|
|
f"cast converter currently doesn't support dtype: {out_dtype}"
|
|
)
|
|
cast_layer = network.add_identity(input_tensor)
|
|
cast_layer.set_output_type(0, out_dtype)
|
|
cast_layer.get_output(0).dtype = out_dtype
|
|
set_layer_name(cast_layer, paddle_op)
|
|
return cast_layer.get_output(0)
|
|
|
|
|
|
@converter_registry.register("pd_op.slice")
|
|
def slice_converter(network, paddle_op, inputs):
|
|
input_tensor = inputs[0]
|
|
axes = paddle_op.attrs()["axes"]
|
|
decrease_axis = paddle_op.attrs().get("decrease_axis")
|
|
|
|
input_shape_tensor = trt_shape(
|
|
network, input_tensor, name=[paddle_op.name(), "input_shape_tensor"]
|
|
)
|
|
input_rank = len(input_tensor.shape)
|
|
|
|
starts_tensor = []
|
|
ends_tensor = []
|
|
for i in range(input_rank):
|
|
starts_tensor.append(
|
|
add_1D_constant_layer(
|
|
network, 0, name=[paddle_op.name(), f'starts_tensor_{i}']
|
|
)
|
|
)
|
|
ends_tensor.append(
|
|
get_shape_tensor_element(
|
|
network,
|
|
input_shape_tensor,
|
|
i,
|
|
name=[paddle_op.name(), f'end_tensor_{i}'],
|
|
)
|
|
)
|
|
|
|
starts = get_input_constant_value(paddle_op, inputs, 1)
|
|
if starts is not None:
|
|
assert len(starts) == len(axes), (
|
|
f"The size of this starts: {len(starts)} must be equal to the axes: {len(axes)}."
|
|
)
|
|
for idx in range(len(axes)):
|
|
if starts[idx] < 0:
|
|
starts_tensor[axes[idx]] = trt_max(
|
|
network,
|
|
trt_sum(
|
|
network,
|
|
add_1D_constant_layer(
|
|
network,
|
|
starts[idx],
|
|
name=[paddle_op.name(), f'starts[idx]_{idx}'],
|
|
),
|
|
get_shape_tensor_element(
|
|
network,
|
|
input_shape_tensor,
|
|
axes[idx],
|
|
name=[paddle_op.name(), f'axes[idx]_{idx}'],
|
|
),
|
|
name=[paddle_op.name(), 'trt_sum'],
|
|
),
|
|
add_1D_constant_layer(
|
|
network, 0, name=[paddle_op.name(), 'zero_tensor']
|
|
),
|
|
name=[
|
|
paddle_op.name(),
|
|
f'starts_tensor[axes[idx]]_{axes[idx]}',
|
|
],
|
|
)
|
|
else:
|
|
starts_tensor[axes[idx]] = trt_min(
|
|
network,
|
|
add_1D_constant_layer(
|
|
network,
|
|
starts[idx],
|
|
name=[paddle_op.name(), f'starts[idx]_{idx}'],
|
|
),
|
|
get_shape_tensor_element(
|
|
network,
|
|
input_shape_tensor,
|
|
axes[idx],
|
|
name=[paddle_op.name(), f'axes[idx]_{idx}'],
|
|
),
|
|
)
|
|
else:
|
|
starts = inputs[1]
|
|
for idx in range(len(axes)):
|
|
starts_tensor[axes[idx]] = get_shape_tensor_element(
|
|
network,
|
|
starts,
|
|
idx,
|
|
name=[paddle_op.name(), f'starts_tensor_{idx}'],
|
|
)
|
|
|
|
ends = get_input_constant_value(paddle_op, inputs, 2)
|
|
if ends is not None:
|
|
assert len(ends) == len(axes), (
|
|
f"The size of this ends: {len(ends)} must be equal to the axes: {len(axes)}."
|
|
)
|
|
for idx in range(len(axes)):
|
|
if ends[idx] < 0:
|
|
ends_tensor[axes[idx]] = trt_max(
|
|
network,
|
|
trt_sum(
|
|
network,
|
|
add_1D_constant_layer(
|
|
network,
|
|
ends[idx],
|
|
name=[paddle_op.name(), f'ends[idx]_{idx}'],
|
|
),
|
|
get_shape_tensor_element(
|
|
network,
|
|
input_shape_tensor,
|
|
axes[idx],
|
|
name=[paddle_op.name(), f'axes[idx]_{idx}'],
|
|
),
|
|
name=[paddle_op.name(), 'trt_sum'],
|
|
),
|
|
add_1D_constant_layer(
|
|
network, 0, name=[paddle_op.name(), 'zero_tensor']
|
|
),
|
|
name=[
|
|
paddle_op.name(),
|
|
f'ends_tensor[axes[idx]]_{axes[idx]}',
|
|
],
|
|
)
|
|
else:
|
|
ends_tensor[axes[idx]] = trt_min(
|
|
network,
|
|
add_1D_constant_layer(
|
|
network,
|
|
ends[idx],
|
|
name=[paddle_op.name(), f'ends[idx]_{idx}'],
|
|
),
|
|
get_shape_tensor_element(
|
|
network,
|
|
input_shape_tensor,
|
|
axes[idx],
|
|
name=[paddle_op.name(), f'axes[idx]_{idx}'],
|
|
),
|
|
)
|
|
else:
|
|
ends = inputs[2]
|
|
for idx in range(len(axes)):
|
|
ends_tensor[axes[idx]] = get_shape_tensor_element(
|
|
network,
|
|
ends,
|
|
idx,
|
|
name=[paddle_op.name(), f'ends_tensor_{idx}'],
|
|
)
|
|
|
|
start_tensor_layer = network.add_concatenation(starts_tensor)
|
|
start_tensor_layer.axis = 0
|
|
set_layer_name(start_tensor_layer, paddle_op)
|
|
start_tensor = start_tensor_layer.get_output(0)
|
|
end_tensor_layer = network.add_concatenation(ends_tensor)
|
|
end_tensor_layer.axis = 0
|
|
set_layer_name(end_tensor_layer, paddle_op)
|
|
end_tensor = end_tensor_layer.get_output(0)
|
|
size_tensor = trt_sub(
|
|
network,
|
|
end_tensor,
|
|
start_tensor,
|
|
name=[paddle_op.name(), 'size_tensor'],
|
|
)
|
|
|
|
# Create Slice layer
|
|
slice_layer = network.add_slice(
|
|
input_tensor, [0] * input_rank, [0] * input_rank, [1] * input_rank
|
|
)
|
|
slice_layer.set_input(1, start_tensor)
|
|
slice_layer.set_input(2, size_tensor)
|
|
set_layer_name(slice_layer, paddle_op)
|
|
|
|
output_tensor = slice_layer.get_output(0)
|
|
|
|
# Handle decrease_axis
|
|
if len(decrease_axis) > 0:
|
|
gather_indices = []
|
|
for i in range(input_rank):
|
|
if i in decrease_axis:
|
|
continue
|
|
gather_indices.append(i)
|
|
|
|
if len(gather_indices) == 0:
|
|
# 0-dim tensor situation and shuffle layer will make its shape (1,) -> ()
|
|
shuffle_layer = network.add_shuffle(output_tensor)
|
|
shuffle_layer.reshape_dims = ()
|
|
else:
|
|
real_size_tensor = trt_gather(network, size_tensor, gather_indices)
|
|
shuffle_layer = network.add_shuffle(output_tensor)
|
|
shuffle_layer.set_input(1, real_size_tensor)
|
|
|
|
set_layer_name(shuffle_layer, paddle_op)
|
|
output_tensor = shuffle_layer.get_output(0)
|
|
|
|
return output_tensor
|
|
|
|
|
|
@converter_registry.register("pd_op.split_with_num")
|
|
def split_with_num_converter(network, paddle_op, inputs):
|
|
input_tensor = inputs[0]
|
|
input_shape_size = len(input_tensor.shape)
|
|
|
|
# Handle the case where axis is of type pir::Value
|
|
axis_value = get_input_constant_value(paddle_op, inputs, 1)
|
|
if axis_value is not None:
|
|
axis_tensor = add_1D_constant_layer(
|
|
network, axis_value, name=[paddle_op.name(), 'axis_tensor']
|
|
)
|
|
else:
|
|
axis_tensor = inputs[1]
|
|
axis_tensor = cast_tensor(
|
|
network,
|
|
axis_tensor,
|
|
trt.int32,
|
|
name=[paddle_op.name(), 'axis_tensor'],
|
|
)
|
|
|
|
num_splits = paddle_op.attrs().get("num")
|
|
num_splits_tensor = add_1D_constant_layer(
|
|
network, num_splits, name=[paddle_op.name(), 'num_splits_tensor']
|
|
)
|
|
|
|
# Get the dynamic shape of the input tensor
|
|
input_shape_tensor = network.add_shape(input_tensor)
|
|
set_layer_name(input_shape_tensor, paddle_op)
|
|
input_shape_tensor = input_shape_tensor.get_output(0)
|
|
|
|
# Handle negative axis index
|
|
input_shape_size_tensor = add_1D_constant_layer(
|
|
network,
|
|
input_shape_size,
|
|
name=[paddle_op.name(), 'input_shape_size_tensor'],
|
|
)
|
|
zero_tensor = add_1D_constant_layer(
|
|
network, 0, name=[paddle_op.name(), 'zero_tensor']
|
|
)
|
|
|
|
is_negative_axis = trt_less(
|
|
network,
|
|
axis_tensor,
|
|
zero_tensor,
|
|
name=[paddle_op.name(), 'is_negative_axis'],
|
|
)
|
|
is_negative_axis_int = cast_tensor(
|
|
network,
|
|
is_negative_axis,
|
|
trt.int32,
|
|
name=[paddle_op.name(), 'is_negative_axis_int'],
|
|
)
|
|
|
|
axis_adjustment = trt_prod(
|
|
network,
|
|
is_negative_axis_int,
|
|
input_shape_size_tensor,
|
|
name=[paddle_op.name(), 'axis_adjustment'],
|
|
)
|
|
|
|
axis_tensor = trt_sum(
|
|
network,
|
|
axis_tensor,
|
|
axis_adjustment,
|
|
name=[paddle_op.name(), 'axis_tensor'],
|
|
)
|
|
|
|
# Get the size of the dimension specified by axis
|
|
input_axis_size = network.add_gather(
|
|
input_shape_tensor, axis_tensor, axis=0
|
|
)
|
|
set_layer_name(input_axis_size, paddle_op)
|
|
input_axis_size = input_axis_size.get_output(0)
|
|
|
|
# Compute the size of each split
|
|
split_size = trt_floor_div(
|
|
network,
|
|
input_axis_size,
|
|
num_splits_tensor,
|
|
name=[paddle_op.name(), 'split_size'],
|
|
)
|
|
|
|
outputs = []
|
|
current_offset = add_1D_constant_layer(
|
|
network, 0, name=[paddle_op.name(), 'current_offset']
|
|
)
|
|
|
|
for idx in range(num_splits):
|
|
idx_tensor = add_1D_constant_layer(
|
|
network, idx, name=[paddle_op.name(), f'idx_tensor_{idx}']
|
|
)
|
|
# Calculate the slice start and size
|
|
start_tensor = build_start_tensor(
|
|
network,
|
|
input_shape_size,
|
|
axis_tensor,
|
|
current_offset,
|
|
name=[paddle_op.name(), f'start_tensor_{idx}'],
|
|
)
|
|
size_tensor = build_size_tensor(
|
|
network,
|
|
input_shape_size,
|
|
axis_tensor,
|
|
split_size,
|
|
input_shape_tensor,
|
|
name=[paddle_op.name(), f'size_tensor_{idx}'],
|
|
)
|
|
|
|
# Create Slice layer
|
|
slice_layer = network.add_slice(
|
|
input_tensor,
|
|
[0] * input_shape_size,
|
|
[0] * input_shape_size,
|
|
[1] * input_shape_size,
|
|
)
|
|
slice_layer.set_input(1, start_tensor)
|
|
slice_layer.set_input(2, size_tensor)
|
|
set_layer_name(slice_layer, paddle_op)
|
|
|
|
outputs.append(slice_layer.get_output(0))
|
|
|
|
# Update current_offset for the next slice
|
|
current_offset = trt_sum(
|
|
network,
|
|
current_offset,
|
|
split_size,
|
|
name=[paddle_op.name(), 'current_offset'],
|
|
)
|
|
|
|
return outputs
|
|
|
|
|
|
@converter_registry.register("pd_op.split")
|
|
def split_converter(network, paddle_op, inputs):
|
|
input_tensor = inputs[0]
|
|
input_shape = input_tensor.shape
|
|
input_shape_size = len(input_shape)
|
|
|
|
axis_value = get_input_constant_value(paddle_op, inputs, 2)
|
|
if axis_value is not None:
|
|
axis_tensor = add_1D_constant_layer(
|
|
network, axis_value, name=[paddle_op.name(), 'axis_tensor']
|
|
)
|
|
else:
|
|
axis_tensor = inputs[2]
|
|
axis_tensor = cast_tensor(
|
|
network,
|
|
axis_tensor,
|
|
trt.int32,
|
|
name=[paddle_op.name(), 'axis_tensor'],
|
|
)
|
|
|
|
# Retrieve and process sections
|
|
sections_value = get_input_constant_value(paddle_op, inputs, 1)
|
|
if sections_value is not None:
|
|
section_list = [int(s) for s in sections_value]
|
|
dynamic_sections = False
|
|
else:
|
|
sections_tensor = inputs[1]
|
|
dynamic_sections = True
|
|
|
|
# Get the dynamic shape of the input tensor
|
|
input_shape_tensor = network.add_shape(input_tensor)
|
|
set_layer_name(input_shape_tensor, paddle_op)
|
|
input_shape_tensor = input_shape_tensor.get_output(0)
|
|
|
|
# Handle negative axis index
|
|
input_shape_size_tensor = add_1D_constant_layer(
|
|
network,
|
|
input_shape_size,
|
|
name=[paddle_op.name(), 'input_shape_size_tensor'],
|
|
)
|
|
zero_tensor = add_1D_constant_layer(
|
|
network, 0, name=[paddle_op.name(), 'zero_tensor']
|
|
)
|
|
|
|
is_negative_axis = trt_less(
|
|
network,
|
|
axis_tensor,
|
|
zero_tensor,
|
|
name=[paddle_op.name(), 'is_negative_axis'],
|
|
)
|
|
is_negative_axis_int = cast_tensor(
|
|
network,
|
|
is_negative_axis,
|
|
trt.int32,
|
|
name=[paddle_op.name(), 'is_negative_axis_int'],
|
|
)
|
|
|
|
axis_adjustment = trt_prod(
|
|
network,
|
|
is_negative_axis_int,
|
|
input_shape_size_tensor,
|
|
name=[paddle_op.name(), 'axis_adjustment'],
|
|
)
|
|
axis_tensor = trt_sum(
|
|
network,
|
|
axis_tensor,
|
|
axis_adjustment,
|
|
name=[paddle_op.name(), 'axis_tensor'],
|
|
)
|
|
|
|
# Initialize output list
|
|
outputs = []
|
|
offset = add_1D_constant_layer(
|
|
network, 0, name=[paddle_op.name(), 'offset']
|
|
)
|
|
|
|
if not dynamic_sections:
|
|
for section_size in section_list:
|
|
section_size_tensor = add_1D_constant_layer(
|
|
network,
|
|
section_size,
|
|
name=[paddle_op.name(), f'section_size_tensor_{section_size}'],
|
|
)
|
|
|
|
# Build start_tensor
|
|
start_tensor = build_start_tensor(
|
|
network,
|
|
input_shape_size,
|
|
axis_tensor,
|
|
offset,
|
|
name=[paddle_op.name(), f'start_tensor_{section_size}'],
|
|
)
|
|
|
|
# Build size_tensor
|
|
size_tensor = build_size_tensor(
|
|
network,
|
|
input_shape_size,
|
|
axis_tensor,
|
|
section_size_tensor,
|
|
input_shape_tensor,
|
|
name=[paddle_op.name(), f'size_tensor_{section_size}'],
|
|
)
|
|
# Create Slice layer
|
|
slice_layer = network.add_slice(
|
|
input_tensor,
|
|
[0] * input_shape_size,
|
|
[0] * input_shape_size,
|
|
[1] * input_shape_size,
|
|
)
|
|
slice_layer.set_input(1, start_tensor)
|
|
slice_layer.set_input(2, size_tensor)
|
|
set_layer_name(slice_layer, paddle_op)
|
|
|
|
outputs.append(slice_layer.get_output(0))
|
|
|
|
# Update offset
|
|
offset = network.add_elementwise(
|
|
offset, section_size_tensor, trt.ElementWiseOperation.SUM
|
|
)
|
|
set_layer_name(offset, paddle_op)
|
|
offset = offset.get_output(0)
|
|
else:
|
|
# If sections is a dynamic tensor
|
|
num_sections = sections_tensor.shape[0]
|
|
if num_sections == -1:
|
|
raise NotImplementedError("dynamic sections not support")
|
|
num_sections = int(num_sections)
|
|
|
|
for idx in range(num_sections):
|
|
idx_tensor = add_1D_constant_layer(
|
|
network, idx, name=[paddle_op.name(), f'idx_tensor_{idx}']
|
|
)
|
|
|
|
# Get section_size_tensor = sections_tensor[idx]
|
|
section_size_tensor = network.add_gather(
|
|
sections_tensor, idx_tensor, axis=0
|
|
)
|
|
set_layer_name(section_size_tensor, paddle_op)
|
|
section_size_tensor = section_size_tensor.get_output(0)
|
|
|
|
# Build start_tensor
|
|
start_tensor = build_start_tensor(
|
|
network,
|
|
input_shape_size,
|
|
axis_tensor,
|
|
offset,
|
|
name=[paddle_op.name(), f'start_tensor_{idx}'],
|
|
)
|
|
|
|
# Build size_tensor
|
|
size_tensor = build_size_tensor(
|
|
network,
|
|
input_shape_size,
|
|
axis_tensor,
|
|
section_size_tensor,
|
|
input_shape_tensor,
|
|
name=[paddle_op.name(), f'size_tensor_{idx}'],
|
|
)
|
|
|
|
# Create Slice layer
|
|
slice_layer = network.add_slice(
|
|
input_tensor,
|
|
[0] * input_shape_size,
|
|
[0] * input_shape_size,
|
|
[1] * input_shape_size,
|
|
)
|
|
slice_layer.set_input(1, start_tensor)
|
|
slice_layer.set_input(2, size_tensor)
|
|
set_layer_name(slice_layer, paddle_op)
|
|
|
|
outputs.append(slice_layer.get_output(0))
|
|
|
|
# Update offset
|
|
offset = network.add_elementwise(
|
|
offset, section_size_tensor, trt.ElementWiseOperation.SUM
|
|
)
|
|
set_layer_name(offset, paddle_op)
|
|
offset = offset.get_output(0)
|
|
|
|
return outputs
|
|
|
|
|
|
@converter_registry.register("pd_op.stack")
|
|
def stack_converter(network, paddle_op, inputs):
|
|
input_tensors = inputs[0]
|
|
input_num = len(input_tensors)
|
|
|
|
inputs = []
|
|
for i in range(input_num):
|
|
inputs.append(input_tensors[i])
|
|
|
|
input_rank = len(input_tensors[0].shape)
|
|
|
|
output_rank = input_rank + 1
|
|
axis = paddle_op.attrs().get("axis")
|
|
if axis < 0:
|
|
axis += output_rank
|
|
|
|
shape_tensor = network.add_shape(input_tensors[0])
|
|
set_layer_name(shape_tensor, paddle_op)
|
|
shape_tensor = shape_tensor.get_output(0)
|
|
shape_tensor_vec = []
|
|
for i in range(output_rank):
|
|
if i < axis:
|
|
shape_tensor_vec.append(
|
|
get_shape_tensor_element(
|
|
network,
|
|
shape_tensor,
|
|
i,
|
|
name=[paddle_op.name(), f'shape_tensor_vec_{i}'],
|
|
)
|
|
)
|
|
elif i > axis:
|
|
shape_tensor_vec.append(
|
|
get_shape_tensor_element(
|
|
network,
|
|
shape_tensor,
|
|
i - 1,
|
|
name=[paddle_op.name(), f'shape_tensor_vec_{i}'],
|
|
)
|
|
)
|
|
else:
|
|
shape_tensor_vec.append(
|
|
add_1D_constant_layer(
|
|
network, 1, name=[paddle_op.name(), f'shape_tensor_vec_{i}']
|
|
)
|
|
)
|
|
|
|
after_shape_tensor = network.add_concatenation(shape_tensor_vec)
|
|
set_layer_name(after_shape_tensor, paddle_op)
|
|
after_shape_tensor = after_shape_tensor.get_output(0)
|
|
|
|
for i in range(input_num):
|
|
shuffle_layer = network.add_shuffle(inputs[i])
|
|
shuffle_layer.set_input(1, after_shape_tensor)
|
|
set_layer_name(shuffle_layer, [paddle_op.name(), f'shuffle_layer_{i}'])
|
|
reshaped_tensor = shuffle_layer.get_output(0)
|
|
inputs[i] = reshaped_tensor
|
|
|
|
concat_layer = network.add_concatenation(inputs)
|
|
concat_layer.axis = axis
|
|
set_layer_name(concat_layer, paddle_op)
|
|
output_tensor = concat_layer.get_output(0)
|
|
|
|
# Because we change tensor to 1-dim in 0-dim tensor situation when use trt,
|
|
# so after stack, output will become 2-dim, if paddle output is a 1d tensor, we need reshape it.
|
|
if (
|
|
len(paddle_op.results()[0].shape) == 1
|
|
and paddle_op.results()[0].shape[0] != -1
|
|
):
|
|
output_tensor = resize_to_1d(
|
|
network, output_tensor, name=[paddle_op.name(), 'output_tensor']
|
|
)
|
|
return output_tensor
|
|
|
|
|
|
@converter_registry.register("pd_op.tile")
|
|
def tile_converter(network, paddle_op, inputs):
|
|
input = inputs[0]
|
|
input_shape = input.shape
|
|
input_shape_tensor = network.add_shape(input)
|
|
set_layer_name(input_shape_tensor, paddle_op)
|
|
input_shape_tensor = input_shape_tensor.get_output(0)
|
|
rank = len(input_shape)
|
|
|
|
repeat_times = get_input_constant_value(paddle_op, inputs, 1)
|
|
if repeat_times is not None:
|
|
repeat_tensor = add_1D_constant_layer(
|
|
network, repeat_times, name=[paddle_op.name(), 'repeat_tensor']
|
|
)
|
|
repeat_rank = len(repeat_times)
|
|
else:
|
|
repeat_tensor = inputs[1]
|
|
if isinstance(repeat_tensor, list):
|
|
repeat_rank = len(repeat_tensor)
|
|
repeat_tensor = trt_concat(
|
|
network, repeat_tensor, name=[paddle_op.name(), 'repeat_tensor']
|
|
)
|
|
else:
|
|
repeat_tensor = resize_to_1d(
|
|
network, repeat_tensor, name=[paddle_op.name(), 'repeat_tensor']
|
|
)
|
|
repeat_shape = paddle_op.operands()[1].source().shape
|
|
repeat_rank = repeat_shape[0]
|
|
|
|
if rank > repeat_rank:
|
|
one_rank_tensor = add_1D_constant_layer(
|
|
network,
|
|
[1] * (rank - repeat_rank),
|
|
name=[paddle_op.name(), 'one_rank_tensor'],
|
|
)
|
|
repeat_expand_tensor = trt_concat(
|
|
network,
|
|
[one_rank_tensor, repeat_tensor],
|
|
name=[paddle_op.name(), 'repeat_expand_tensor'],
|
|
)
|
|
elif rank < repeat_rank:
|
|
one_rank_tensor = add_1D_constant_layer(
|
|
network,
|
|
[1] * (repeat_rank - rank),
|
|
name=[paddle_op.name(), 'one_rank_tensor'],
|
|
)
|
|
input_shape_tensor = trt_concat(
|
|
network,
|
|
[one_rank_tensor, input_shape_tensor],
|
|
name=[paddle_op.name(), 'input_shape_tensor'],
|
|
)
|
|
input = trt_reshape(
|
|
network,
|
|
input,
|
|
input_shape_tensor,
|
|
name=[paddle_op.name(), 'input_shape_tensor'],
|
|
is_shape_tensor=True,
|
|
)
|
|
repeat_expand_tensor = repeat_tensor
|
|
else:
|
|
repeat_expand_tensor = repeat_tensor
|
|
|
|
start = [0] * max(rank, repeat_rank)
|
|
stride = [1] * max(rank, repeat_rank)
|
|
output_shape = [0] * max(rank, repeat_rank)
|
|
output_shape_tensor = trt_prod(
|
|
network,
|
|
input_shape_tensor,
|
|
repeat_expand_tensor,
|
|
name=[paddle_op.name(), 'output_shape_tensor'],
|
|
)
|
|
|
|
slice_layer = network.add_slice(input, start, output_shape, stride)
|
|
slice_layer.set_input(2, output_shape_tensor)
|
|
set_layer_name(slice_layer, paddle_op)
|
|
|
|
version_list = get_trt_version_list()
|
|
if version_list >= [8, 6, 0]:
|
|
slice_layer.mode = trt.SampleMode.WRAP
|
|
else:
|
|
slice_layer.mode = trt.SliceMode.WRAP
|
|
|
|
return slice_layer.get_output(0)
|
|
|
|
|
|
@converter_registry.register(
|
|
"pd_op.take_along_axis", trt_version="trt_version_ge=8.2"
|
|
)
|
|
def take_along_axis_converter(network, paddle_op, inputs):
|
|
axis = paddle_op.attrs().get("axis", 0)
|
|
input_tensor = inputs[0]
|
|
index_tensor = inputs[1]
|
|
|
|
input_dims = input_tensor.shape
|
|
if axis < 0:
|
|
axis += len(input_dims)
|
|
|
|
gather_layer = network.add_gather_v2(
|
|
input_tensor, index_tensor, trt.GatherMode.ELEMENT
|
|
)
|
|
gather_layer.axis = axis
|
|
set_layer_name(gather_layer, paddle_op)
|
|
|
|
output_tensor = gather_layer.get_output(0)
|
|
|
|
return output_tensor
|
|
|
|
|
|
@converter_registry.register("pd_op.strided_slice")
|
|
def strided_slice_converter(network, paddle_op, inputs):
|
|
input_tensor = inputs[0]
|
|
axes = paddle_op.attrs()["axes"]
|
|
starts = get_input_constant_value(paddle_op, inputs, 1)
|
|
ends = get_input_constant_value(paddle_op, inputs, 2)
|
|
strides = get_input_constant_value(paddle_op, inputs, 3)
|
|
|
|
input_shape = input_tensor.shape
|
|
nchw_input_dims = len(input_shape)
|
|
|
|
trt_start_dims = [0] * nchw_input_dims
|
|
trt_end_dims = [0] * nchw_input_dims
|
|
trt_size_dims = [0] * nchw_input_dims
|
|
trt_step_dims = [1] * nchw_input_dims
|
|
|
|
has_neg_indices = False
|
|
|
|
for i, trt_axis in enumerate(axes):
|
|
trt_start_dims[trt_axis] = starts[i]
|
|
trt_end_dims[trt_axis] = ends[i]
|
|
trt_step_dims[trt_axis] = strides[i]
|
|
if starts[i] < 0 or ends[i] < 0:
|
|
has_neg_indices = True
|
|
|
|
shape_tensor = trt_shape(
|
|
network, input_tensor, name=[paddle_op.name(), 'shape_tensor']
|
|
)
|
|
start_tensor = add_1D_constant_layer(
|
|
network, trt_start_dims, name=[paddle_op.name(), 'start_tensor']
|
|
)
|
|
if has_neg_indices:
|
|
start_tensor = fix_negative_indices(
|
|
network,
|
|
shape_tensor,
|
|
start_tensor,
|
|
name=[paddle_op.name(), 'start_tensor'],
|
|
)
|
|
|
|
end_vec_tensor = []
|
|
for i in range(len(trt_end_dims)):
|
|
end_vec_tensor.append(
|
|
get_shape_tensor_element(
|
|
network,
|
|
shape_tensor,
|
|
i,
|
|
name=[paddle_op.name(), f'end_vec_tensor{i}'],
|
|
)
|
|
)
|
|
|
|
for i, trt_axis in enumerate(axes):
|
|
if ends[i] >= 0:
|
|
end_vec_tensor[trt_axis] = add_1D_constant_layer(
|
|
network, ends[i], name=[paddle_op.name(), f'end_vec_tensor{i}']
|
|
)
|
|
else:
|
|
end_vec_tensor[trt_axis] = trt_sum(
|
|
network,
|
|
end_vec_tensor[trt_axis],
|
|
add_1D_constant_layer(
|
|
network,
|
|
ends[i],
|
|
name=[paddle_op.name(), f'end_vec_tensor{i}'],
|
|
),
|
|
name=[paddle_op.name(), f'end_vec_tensor{i}'],
|
|
)
|
|
|
|
size_tensor = trt_sub(
|
|
network,
|
|
start_tensor,
|
|
trt_min(
|
|
network,
|
|
trt_concat(
|
|
network, end_vec_tensor, name=[paddle_op.name(), 'trt_concat']
|
|
),
|
|
shape_tensor,
|
|
name=[paddle_op.name(), 'trt_min'],
|
|
),
|
|
name=[paddle_op.name(), 'size_tensor'],
|
|
)
|
|
zero_t = add_1D_constant_layer(
|
|
network, 0, name=[paddle_op.name(), 'zero_t']
|
|
)
|
|
step_tensor = add_1D_constant_layer(
|
|
network, trt_step_dims, name=[paddle_op.name(), 'step_tensor']
|
|
)
|
|
size_tensor = trt_sub(
|
|
network,
|
|
zero_t,
|
|
trt_floor_div(
|
|
network,
|
|
size_tensor,
|
|
step_tensor,
|
|
name=[paddle_op.name(), 'trt_floor_div'],
|
|
),
|
|
name=[paddle_op.name(), 'size_tensor'],
|
|
)
|
|
|
|
layer = network.add_slice(
|
|
input_tensor, trt_start_dims, trt_size_dims, trt_step_dims
|
|
)
|
|
layer.set_input(1, start_tensor)
|
|
layer.set_input(2, size_tensor)
|
|
layer.set_input(3, step_tensor)
|
|
set_layer_name(layer, paddle_op)
|
|
return layer.get_output(0)
|
|
|
|
|
|
@converter_registry.register("pd_op.roll")
|
|
def roll_converter(network, paddle_op, inputs):
|
|
input_tensor = inputs[0]
|
|
axis = paddle_op.attrs()["axis"]
|
|
|
|
shifts = get_input_constant_value(paddle_op, inputs, 1)
|
|
if shifts is None:
|
|
shifts = inputs[1]
|
|
|
|
axis_size = len(axis)
|
|
input_shape_tensor = trt_shape(
|
|
network, input_tensor, name=[paddle_op.name(), 'input_shape_tensor']
|
|
)
|
|
|
|
for i in range(axis_size):
|
|
axi = axis[i]
|
|
if isinstance(shifts, trt.ITensor):
|
|
shift = get_shape_tensor_element(
|
|
network, shifts, i, name=[paddle_op.name(), f'shift_{i}']
|
|
)
|
|
input_shift = shift
|
|
else:
|
|
shift = shifts[i]
|
|
input_shift = add_1D_constant_layer(
|
|
network, shift, name=[paddle_op.name(), f'input_shift_{i}']
|
|
)
|
|
input_axis = get_shape_tensor_element(
|
|
network,
|
|
input_shape_tensor,
|
|
axi,
|
|
name=[paddle_op.name(), f'input_axis_{i}'],
|
|
)
|
|
|
|
# 1.sub_value mod input_axis
|
|
input1 = trt_sub(
|
|
network,
|
|
input_axis,
|
|
input_shift,
|
|
name=[paddle_op.name(), f'input1_{i}'],
|
|
)
|
|
tmp_div_res = trt_floor_div(
|
|
network,
|
|
input1,
|
|
input_axis,
|
|
name=[paddle_op.name(), f'tmp_div_res_{i}'],
|
|
)
|
|
tmp_prod_res = trt_prod(
|
|
network,
|
|
tmp_div_res,
|
|
input_axis,
|
|
name=[paddle_op.name(), f'tmp_prod_res_{i}'],
|
|
)
|
|
start = trt_sub(
|
|
network, input1, tmp_prod_res, name=[paddle_op.name(), f'start_{i}']
|
|
)
|
|
# 2.avoid start less than 0,start mod input_axis
|
|
start = trt_sum(
|
|
network, start, input_axis, name=[paddle_op.name(), f'start_{i}']
|
|
)
|
|
tmp_div_res1 = trt_floor_div(
|
|
network,
|
|
start,
|
|
input_axis,
|
|
name=[paddle_op.name(), f'tmp_div_res1_{i}'],
|
|
)
|
|
tmp_prod_res1 = trt_prod(
|
|
network,
|
|
tmp_div_res1,
|
|
input_axis,
|
|
name=[paddle_op.name(), f'tmp_prod_res1_{i}'],
|
|
)
|
|
start = trt_sub(
|
|
network, start, tmp_prod_res1, name=[paddle_op.name(), f'start_{i}']
|
|
)
|
|
zero_tensor = add_1D_constant_layer(
|
|
network, 0, name=[paddle_op.name(), f'zero_tensor_{i}']
|
|
)
|
|
step = add_1D_constant_layer(
|
|
network, 1, name=[paddle_op.name(), f'step_{i}']
|
|
)
|
|
# 3.make index_tensor0
|
|
sub_qutient = trt_sub(
|
|
network,
|
|
input_axis,
|
|
start,
|
|
name=[paddle_op.name(), f'sub_qutient_{i}'],
|
|
)
|
|
quotient_tensor = trt_floor_div(
|
|
network,
|
|
sub_qutient,
|
|
step,
|
|
name=[paddle_op.name(), f'quotient_tensor_{i}'],
|
|
)
|
|
start1 = get_shape_tensor_element(
|
|
network,
|
|
start,
|
|
0,
|
|
is_scalar=True,
|
|
name=[paddle_op.name(), f'start1_{i}'],
|
|
)
|
|
fill_layer0 = network.add_fill(shape=(), op=trt.FillOperation.LINSPACE)
|
|
fill_layer0.set_input(0, quotient_tensor)
|
|
fill_layer0.set_input(1, start1)
|
|
fill_layer0.set_input(2, step)
|
|
set_layer_name(fill_layer0, paddle_op)
|
|
index_tensor0 = fill_layer0.get_output(0)
|
|
# 4.make index_tensor1
|
|
sub_qutient_tensor = trt_sub(
|
|
network,
|
|
start,
|
|
zero_tensor,
|
|
name=[paddle_op.name(), f'sub_qutient_tensor_{i}'],
|
|
)
|
|
quotient_tensor = trt_floor_div(
|
|
network,
|
|
sub_qutient_tensor,
|
|
step,
|
|
name=[paddle_op.name(), f'quotient_tensor_{i}'],
|
|
)
|
|
start2 = add_1D_constant_layer(
|
|
network, 0, is_scalar=True, name=[paddle_op.name(), f'start2_{i}']
|
|
)
|
|
fill_layer1 = network.add_fill(shape=(), op=trt.FillOperation.LINSPACE)
|
|
fill_layer1.set_input(0, quotient_tensor)
|
|
fill_layer1.set_input(1, start2)
|
|
fill_layer1.set_input(2, step)
|
|
set_layer_name(fill_layer1, paddle_op)
|
|
index_tensor1 = fill_layer1.get_output(0)
|
|
itensors = [index_tensor0, index_tensor1]
|
|
concat_input_tensor = trt_concat(
|
|
network, itensors, name=[paddle_op.name(), 'concat_input_tensor']
|
|
)
|
|
if i == 0:
|
|
layer = network.add_gather(
|
|
input=input_tensor, indices=concat_input_tensor, axis=axi
|
|
)
|
|
else:
|
|
layer = network.add_gather(
|
|
input=layer.get_output(0), indices=concat_input_tensor, axis=axi
|
|
)
|
|
set_layer_name(layer, paddle_op)
|
|
|
|
return layer.get_output(0)
|
|
|
|
|
|
@converter_registry.register("pd_op.pad")
|
|
def pad_converter(network, paddle_op, inputs):
|
|
input_tensor = inputs[0]
|
|
paddings = paddle_op.attrs()["paddings"]
|
|
pad_size = len(paddings)
|
|
pre_pad = [paddings[pad_size - 4], paddings[pad_size - 2]]
|
|
post_pad = [paddings[pad_size - 3], paddings[pad_size - 1]]
|
|
layer = network.add_padding_nd(input_tensor, pre_pad, post_pad)
|
|
set_layer_name(layer, paddle_op)
|
|
return layer.get_output(0)
|
|
|
|
|
|
@converter_registry.register("pd_op.pad3d")
|
|
def pad3d_converter(network, paddle_op, inputs):
|
|
input_tensor, paddings = inputs
|
|
value = paddle_op.attrs().get("pad_value", 0.0)
|
|
padding_mode = paddle_op.attrs().get("mode", "constant")
|
|
data_format = paddle_op.attrs().get("data_format")
|
|
if padding_mode == "circular" or data_format == "NDHWC":
|
|
attrs = paddle_op.attrs()
|
|
value_attr = get_input_constant_value(paddle_op, inputs, 1)
|
|
attrs["paddings"] = value_attr
|
|
layer = generic_plugin_converter(network, paddle_op, inputs, attrs)
|
|
return layer.get_output(0)
|
|
else:
|
|
input_dim = len(input_tensor.shape)
|
|
pad_size = paddings.shape[0]
|
|
assert input_dim * 2 - 4 == pad_size, (
|
|
f"Expected paddings size is {input_dim * 2 - 4}, but received {pad_size}."
|
|
)
|
|
|
|
shuffle_index = [4, 2, 0, 5, 3, 1]
|
|
shuffle_inputs = [
|
|
get_shape_tensor_element(
|
|
network,
|
|
paddings,
|
|
shuffle_index[i],
|
|
name=[paddle_op.name(), f'shuffle_inputs_{i}'],
|
|
)
|
|
for i in range(pad_size)
|
|
]
|
|
paddings = trt_concat(
|
|
network, shuffle_inputs, name=[paddle_op.name(), 'paddings']
|
|
)
|
|
|
|
pre_zeros = add_1D_constant_layer(
|
|
network, [0, 0], name=[paddle_op.name(), 'pre_zeros']
|
|
)
|
|
start_slice1 = [0]
|
|
start_slice2 = [3]
|
|
size_slice = [3]
|
|
stride_slice = [1]
|
|
pre_pad = network.add_slice(
|
|
paddings, start_slice1, size_slice, stride_slice
|
|
)
|
|
set_layer_name(pre_pad, paddle_op)
|
|
pre_pad = pre_pad.get_output(0)
|
|
pre_pad = trt_concat(
|
|
network, [pre_zeros, pre_pad], name=[paddle_op.name(), 'pre_pad']
|
|
)
|
|
post_pad = network.add_slice(
|
|
paddings, start_slice2, size_slice, stride_slice
|
|
)
|
|
set_layer_name(post_pad, paddle_op)
|
|
post_pad = post_pad.get_output(0)
|
|
post_pad = trt_concat(
|
|
network, [pre_zeros, post_pad], name=[paddle_op.name(), 'post_pad']
|
|
)
|
|
|
|
zeros = add_1D_constant_layer(
|
|
network, [0] * input_dim, 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']
|
|
)
|
|
|
|
# Add slice layer
|
|
stride = [1] * input_dim
|
|
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)
|
|
|
|
# Set padding mode
|
|
if padding_mode == "constant":
|
|
slice_layer.mode = trt.SampleMode.FILL
|
|
if value != 0.0:
|
|
if input_tensor.dtype in (
|
|
trt.DataType.FLOAT,
|
|
trt.DataType.HALF,
|
|
trt.DataType.INT8,
|
|
):
|
|
fill_value = add_1D_constant_layer(
|
|
network,
|
|
value,
|
|
dtype=np.float32,
|
|
name=[paddle_op.name(), 'fill_value'],
|
|
)
|
|
else:
|
|
value_int = int(value)
|
|
fill_value = add_1D_constant_layer(
|
|
network,
|
|
value_int,
|
|
dtype=np.int32,
|
|
name=[paddle_op.name(), 'fill_value'],
|
|
)
|
|
slice_layer.set_input(4, fill_value)
|
|
elif padding_mode == "reflect":
|
|
slice_layer.mode = trt.SampleMode.REFLECT
|
|
elif padding_mode == "replicate":
|
|
slice_layer.mode = trt.SampleMode.CLAMP
|
|
else:
|
|
raise ValueError(f"Unsupported padding mode: {padding_mode}")
|
|
|
|
return slice_layer.get_output(0)
|
|
|
|
|
|
@converter_registry.register("pd_op.numel")
|
|
def numel_converter(network, paddle_op, inputs):
|
|
input_tensor = inputs[0]
|
|
shape_tensor = network.add_shape(input_tensor)
|
|
set_layer_name(shape_tensor, paddle_op)
|
|
shape_tensor = shape_tensor.get_output(0)
|
|
layer = network.add_reduce(
|
|
shape_tensor, trt.ReduceOperation.PROD, axes=1, keep_dims=False
|
|
)
|
|
set_layer_name(layer, paddle_op)
|
|
return layer.get_output(0)
|
|
|
|
|
|
@converter_registry.register("pd_op.index_put")
|
|
def index_put_converter(network, paddle_op, inputs):
|
|
input_tensor, indices_list, value_tensor = inputs
|
|
indices_tensor = indices_list[0]
|
|
input_shape_tensor = trt_shape(
|
|
network, input_tensor, name=[paddle_op.name(), 'input_shape_tensor']
|
|
)
|
|
input_dims = input_tensor.shape
|
|
indices_dims = indices_tensor.shape
|
|
rank = len(input_dims)
|
|
|
|
# indices
|
|
indices_shape_vec = [
|
|
add_1D_constant_layer(
|
|
network,
|
|
indices_dims[i] if i < len(indices_dims) else 1,
|
|
name=[paddle_op.name(), f'indices_shape_vec_{i}'],
|
|
)
|
|
for i in range(rank)
|
|
]
|
|
start_tensor_vec = [
|
|
add_1D_constant_layer(
|
|
network, 0, name=[paddle_op.name(), f'start_tensor_vec_{i}']
|
|
)
|
|
for i in range(rank)
|
|
]
|
|
stride_tensor_vec = [
|
|
add_1D_constant_layer(
|
|
network, 1, name=[paddle_op.name(), f'stride_tensor_vec_{i}']
|
|
)
|
|
for i in range(rank)
|
|
]
|
|
indices_tensor_temp = trt_reshape(
|
|
network,
|
|
indices_tensor,
|
|
trt_concat(
|
|
network,
|
|
indices_shape_vec,
|
|
name=[paddle_op.name(), 'indices_shape_vec'],
|
|
),
|
|
name=[paddle_op.name(), 'indices_tensor_temp'],
|
|
is_shape_tensor=True,
|
|
)
|
|
start_tensor = trt_concat(
|
|
network, start_tensor_vec, name=[paddle_op.name(), 'start_tensor']
|
|
)
|
|
stride_tensor = trt_concat(
|
|
network, stride_tensor_vec, name=[paddle_op.name(), 'stride_tensor']
|
|
)
|
|
|
|
# slice
|
|
stride = [1] * rank
|
|
indices_slice_layer = network.add_slice(
|
|
trt_cast(
|
|
network,
|
|
indices_tensor_temp,
|
|
trt.float32,
|
|
name=[paddle_op.name(), 'indices_tensor_temp'],
|
|
),
|
|
stride,
|
|
stride,
|
|
stride,
|
|
)
|
|
indices_slice_layer.set_input(1, start_tensor)
|
|
indices_slice_layer.set_input(2, input_shape_tensor)
|
|
indices_slice_layer.set_input(3, stride_tensor)
|
|
indices_slice_layer.mode = trt.SampleMode.CLAMP
|
|
set_layer_name(indices_slice_layer, paddle_op)
|
|
|
|
bool_indices_tensor = trt_cast(
|
|
network,
|
|
indices_slice_layer.get_output(0),
|
|
trt.bool,
|
|
name=[paddle_op.name(), 'bool_indices_tensor'],
|
|
)
|
|
|
|
# nonzero
|
|
nonzero_layer = network.add_non_zero(bool_indices_tensor)
|
|
set_layer_name(nonzero_layer, paddle_op)
|
|
indices_tensor = nonzero_layer.get_output(0)
|
|
permutation = trt.Permutation([1, 0])
|
|
trans_layer = network.add_shuffle(indices_tensor)
|
|
trans_layer.first_transpose = permutation
|
|
set_layer_name(trans_layer, paddle_op)
|
|
indices_tensor = trans_layer.get_output(0)
|
|
indices_new_shape_tensor = trt_shape(
|
|
network,
|
|
indices_tensor,
|
|
name=[paddle_op.name(), 'indices_new_shape_tensor'],
|
|
)
|
|
indices_count_tensor = get_shape_tensor_element(
|
|
network,
|
|
indices_new_shape_tensor,
|
|
0,
|
|
name=[paddle_op.name(), 'indices_count_tensor'],
|
|
)
|
|
|
|
# value
|
|
value_stride = [1]
|
|
value_slice_layer = network.add_slice(
|
|
value_tensor, value_stride, value_stride, value_stride
|
|
)
|
|
value_slice_layer.set_input(
|
|
1,
|
|
add_1D_constant_layer(
|
|
network, 0, name=[paddle_op.name(), 'value_slice_layer_start']
|
|
),
|
|
)
|
|
value_slice_layer.set_input(2, indices_count_tensor)
|
|
value_slice_layer.set_input(
|
|
3,
|
|
add_1D_constant_layer(
|
|
network, 1, name=[paddle_op.name(), 'value_slice_layer_stride']
|
|
),
|
|
)
|
|
value_slice_layer.mode = trt.SampleMode.CLAMP
|
|
set_layer_name(value_slice_layer, paddle_op)
|
|
value_tensor = value_slice_layer.get_output(0)
|
|
|
|
layer = network.add_scatter(
|
|
input_tensor, indices_tensor, value_tensor, trt.ScatterMode.ND
|
|
)
|
|
set_layer_name(layer, paddle_op)
|
|
return layer.get_output(0)
|