chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,285 @@
|
||||
# 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 tensorrt as trt
|
||||
|
||||
from paddle.tensorrt.converter_utils import (
|
||||
generic_plugin_converter,
|
||||
get_input_constant_value,
|
||||
get_shape_tensor_element,
|
||||
set_layer_name,
|
||||
squeeze_trt,
|
||||
trt_cast,
|
||||
trt_gather,
|
||||
trt_reshape,
|
||||
trt_shape,
|
||||
trt_unsqueeze,
|
||||
)
|
||||
from paddle.tensorrt.register import converter_registry
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.nonzero")
|
||||
def non_zero_converter(network, paddle_op, inputs):
|
||||
input_tensor = inputs[0]
|
||||
cast_layer = network.add_cast(input_tensor, trt.float32)
|
||||
set_layer_name(cast_layer, paddle_op)
|
||||
|
||||
non_zero_layer = network.add_non_zero(cast_layer.get_output(0))
|
||||
nonzero_output = non_zero_layer.get_output(0)
|
||||
set_layer_name(non_zero_layer, paddle_op)
|
||||
|
||||
shuffle_layer = network.add_shuffle(input=nonzero_output)
|
||||
shuffle_layer.first_transpose = (1, 0)
|
||||
transposed_output = shuffle_layer.get_output(0)
|
||||
set_layer_name(shuffle_layer, paddle_op)
|
||||
return transposed_output
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.argmax")
|
||||
def argmax_converter(network, paddle_op, inputs):
|
||||
x = inputs[0]
|
||||
input_dims = x.shape
|
||||
rank = len(input_dims)
|
||||
axis = int(get_input_constant_value(paddle_op, inputs, 1)[0])
|
||||
keepdims = paddle_op.attrs()["keepdims"]
|
||||
|
||||
if axis < 0:
|
||||
axis += rank
|
||||
|
||||
topk_layer = network.add_topk(
|
||||
input=x, op=trt.TopKOperation.MAX, k=1, axes=(1 << axis)
|
||||
)
|
||||
set_layer_name(topk_layer, paddle_op)
|
||||
|
||||
if keepdims:
|
||||
return topk_layer.get_output(1)
|
||||
else:
|
||||
topk_out = topk_layer.get_output(1)
|
||||
topk_out_shape_size = len(topk_out.shape)
|
||||
# Mark which dimensions to squeeze
|
||||
should_squeeze = [False] * topk_out_shape_size
|
||||
should_squeeze[axis] = True
|
||||
|
||||
# Get dimensions to keep
|
||||
gather_indices = [
|
||||
i for i, squeeze in enumerate(should_squeeze) if not squeeze
|
||||
]
|
||||
|
||||
# Add Shuffle layer
|
||||
layer = network.add_shuffle(topk_out)
|
||||
shape_tensor = trt_shape(
|
||||
network, topk_out, name=[paddle_op.name(), 'shape_tensor']
|
||||
)
|
||||
real_shape_tensor = trt_gather(
|
||||
network,
|
||||
shape_tensor,
|
||||
gather_indices,
|
||||
name=[paddle_op.name(), 'real_shape_tensor'],
|
||||
)
|
||||
layer.set_input(1, real_shape_tensor)
|
||||
set_layer_name(layer, paddle_op)
|
||||
return layer.get_output(0)
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.argmin")
|
||||
def argmin_converter(network, paddle_op, inputs):
|
||||
x = inputs[0]
|
||||
input_dims = x.shape
|
||||
rank = len(input_dims)
|
||||
axis = int(get_input_constant_value(paddle_op, inputs, 1)[0])
|
||||
keepdims = paddle_op.attrs()["keepdims"]
|
||||
|
||||
if axis < 0:
|
||||
axis += rank
|
||||
|
||||
topk_layer = network.add_topk(
|
||||
input=x, op=trt.TopKOperation.MIN, k=1, axes=(1 << axis)
|
||||
)
|
||||
set_layer_name(topk_layer, paddle_op)
|
||||
|
||||
if keepdims:
|
||||
return topk_layer.get_output(1)
|
||||
else:
|
||||
squeeze_layer = network.add_shuffle(topk_layer.get_output(1))
|
||||
set_layer_name(squeeze_layer, paddle_op)
|
||||
output_dims = []
|
||||
for i in range(len(input_dims)):
|
||||
if i == axis:
|
||||
continue
|
||||
output_dims.append(input_dims[i])
|
||||
squeeze_layer.reshape_dims = tuple(output_dims)
|
||||
return squeeze_layer.get_output(0)
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.argsort")
|
||||
def argsort_converter(network, paddle_op, inputs):
|
||||
input_tensor = inputs[0]
|
||||
input_shape = input_tensor.shape
|
||||
in_type = input_tensor.dtype
|
||||
in_rank = len(input_shape)
|
||||
axis = paddle_op.attrs()["axis"]
|
||||
descending = paddle_op.attrs()["descending"]
|
||||
if input_shape[axis] > 3840:
|
||||
layer = generic_plugin_converter(network, paddle_op, inputs)
|
||||
out0 = layer.get_output(0)
|
||||
out1 = layer.get_output(1)
|
||||
return out0, out1
|
||||
else:
|
||||
if axis < 0:
|
||||
axis += len(input_shape)
|
||||
topk_op = trt.TopKOperation.MAX if descending else trt.TopKOperation.MIN
|
||||
need_cast = True if in_type != trt.DataType.FLOAT else False
|
||||
if in_rank == 1:
|
||||
unsqueeze_shape = trt.Dims([1, -1])
|
||||
input_tensor = trt_reshape(
|
||||
network,
|
||||
input_tensor,
|
||||
unsqueeze_shape,
|
||||
is_shape_tensor=False,
|
||||
name=[paddle_op.name(), 'input_tensor'],
|
||||
)
|
||||
axis = 1
|
||||
if need_cast:
|
||||
input_tensor = trt_cast(
|
||||
network,
|
||||
input_tensor,
|
||||
trt.DataType.FLOAT,
|
||||
name=[paddle_op.name(), 'input_tensor'],
|
||||
)
|
||||
topk_layer = network.add_topk(input_tensor, topk_op, 1, 1 << axis)
|
||||
shape = trt_shape(
|
||||
network, input_tensor, name=[paddle_op.name(), 'shape']
|
||||
)
|
||||
k_tensor = get_shape_tensor_element(
|
||||
network, shape, axis, True, name=[paddle_op.name(), 'k_tensor']
|
||||
)
|
||||
topk_layer.set_input(1, k_tensor)
|
||||
set_layer_name(topk_layer, paddle_op)
|
||||
out = topk_layer.get_output(0)
|
||||
indices = topk_layer.get_output(1)
|
||||
if in_rank == 1:
|
||||
squeeze_shape = trt.Dims([-1])
|
||||
out = trt_reshape(
|
||||
network,
|
||||
out,
|
||||
squeeze_shape,
|
||||
is_shape_tensor=False,
|
||||
name=[paddle_op.name(), 'out'],
|
||||
)
|
||||
indices = trt_reshape(
|
||||
network,
|
||||
indices,
|
||||
squeeze_shape,
|
||||
is_shape_tensor=False,
|
||||
name=[paddle_op.name(), 'indices'],
|
||||
)
|
||||
out_tensor = trt_cast(
|
||||
network, out, in_type, name=[paddle_op.name(), 'out_tensor']
|
||||
)
|
||||
indices_tensor = trt_cast(
|
||||
network,
|
||||
indices,
|
||||
indices.dtype,
|
||||
name=[paddle_op.name(), 'indices_tensor'],
|
||||
)
|
||||
return out_tensor, indices_tensor
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.where")
|
||||
def where_converter(network, paddle_op, inputs):
|
||||
condition = inputs[0]
|
||||
x = inputs[1]
|
||||
y = inputs[2]
|
||||
|
||||
select_layer = network.add_select(condition, x, y)
|
||||
set_layer_name(select_layer, paddle_op)
|
||||
|
||||
return select_layer.get_output(0)
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.topk")
|
||||
def topk_converter(network, paddle_op, inputs):
|
||||
input_tensor = inputs[0]
|
||||
|
||||
input_shape = input_tensor.shape
|
||||
|
||||
axis = paddle_op.attrs().get("axis", -1)
|
||||
largest = paddle_op.attrs().get("largest", True)
|
||||
flag = trt.TopKOperation.MAX if largest else trt.TopKOperation.MIN
|
||||
|
||||
k_list = get_input_constant_value(paddle_op, inputs, 1)
|
||||
if k_list is None:
|
||||
raise NotImplementedError("Dynamic k is not supported in TensorRT.")
|
||||
k = k_list[0]
|
||||
input_rank = len(input_shape)
|
||||
|
||||
expand_to_2d = input_rank == 1
|
||||
if expand_to_2d:
|
||||
input_tensor = trt_unsqueeze(
|
||||
network, input_tensor, [1], name=[paddle_op.name(), 'input_tensor']
|
||||
)
|
||||
|
||||
input_type = input_tensor.dtype
|
||||
if input_type == trt.DataType.INT32:
|
||||
input_tensor = trt_cast(
|
||||
network,
|
||||
input_tensor,
|
||||
trt.DataType.FLOAT,
|
||||
name=[paddle_op.name(), 'input_tensor'],
|
||||
)
|
||||
|
||||
if axis < 0:
|
||||
axis += input_rank
|
||||
|
||||
layer = network.add_topk(input_tensor, flag, int(k), 1 << axis)
|
||||
set_layer_name(layer, paddle_op)
|
||||
values = layer.get_output(0)
|
||||
indices = layer.get_output(1)
|
||||
|
||||
if expand_to_2d:
|
||||
values = squeeze_trt(
|
||||
network, values, [1], name=[paddle_op.name(), 'values']
|
||||
)
|
||||
indices = squeeze_trt(
|
||||
network, indices, [1], name=[paddle_op.name(), 'indices']
|
||||
)
|
||||
|
||||
if input_type == trt.DataType.INT32:
|
||||
values = trt_cast(
|
||||
network,
|
||||
values,
|
||||
trt.DataType.INT32,
|
||||
name=[paddle_op.name(), 'values'],
|
||||
)
|
||||
|
||||
return values, indices
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.index_select")
|
||||
def index_select_converter(network, paddle_op, inputs):
|
||||
input_tensor = inputs[0]
|
||||
index_tensor = inputs[1]
|
||||
axis = paddle_op.attrs().get("axis", 0)
|
||||
|
||||
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
|
||||
)
|
||||
set_layer_name(gather_layer, paddle_op)
|
||||
|
||||
return gather_layer.get_output(0)
|
||||
Reference in New Issue
Block a user