chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,556 @@
|
||||
# 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 import pir
|
||||
from paddle.tensorrt.converter_utils import (
|
||||
add_1D_constant_layer,
|
||||
get_input_constant_value,
|
||||
get_shape_tensor_element,
|
||||
set_layer_name,
|
||||
trt_concat,
|
||||
trt_reshape,
|
||||
trt_shape,
|
||||
)
|
||||
from paddle.tensorrt.register import converter_registry
|
||||
from paddle.tensorrt.util import get_trt_version_list
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.dropout")
|
||||
def dropout_converter(network, paddle_op, inputs):
|
||||
input_x = inputs[0]
|
||||
dropout_prob = get_input_constant_value(paddle_op, inputs, 2)[0]
|
||||
downgrade_in_infer = paddle_op.attrs().get("mode")
|
||||
|
||||
if downgrade_in_infer == "upscale_in_train":
|
||||
shuffle_layer = network.add_shuffle(input_x)
|
||||
set_layer_name(shuffle_layer, paddle_op)
|
||||
return shuffle_layer.get_output(0)
|
||||
|
||||
weight_data = np.array([1 - dropout_prob]).astype("float32")
|
||||
scale_weights = trt.Weights(weight_data)
|
||||
shift_weights = trt.Weights(np.array([0]).astype("float32"))
|
||||
power_weights = trt.Weights(np.array([1]).astype("float32"))
|
||||
|
||||
scale_layer = network.add_scale(
|
||||
input_x,
|
||||
mode=trt.ScaleMode.UNIFORM,
|
||||
shift=shift_weights,
|
||||
scale=scale_weights,
|
||||
power=power_weights,
|
||||
)
|
||||
set_layer_name(scale_layer, paddle_op)
|
||||
|
||||
return scale_layer.get_output(0)
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.bilinear_interp")
|
||||
def bilinear_interp_converter(network, paddle_op, inputs):
|
||||
input_tensor = inputs[0]
|
||||
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)
|
||||
|
||||
input_rank = input_shape_tensor.shape # The reason is unknown that adding this unused code make input_shape_tensor maintain the correct result.
|
||||
data_format = paddle_op.attrs().get("data_format")
|
||||
interp_method = paddle_op.attrs().get("interp_method")
|
||||
align_corners = paddle_op.attrs().get("align_corners")
|
||||
align_mode = paddle_op.attrs().get("align_mode")
|
||||
out_h = paddle_op.attrs().get("out_h")
|
||||
out_w = paddle_op.attrs().get("out_w")
|
||||
out_d = paddle_op.attrs().get("out_d")
|
||||
scale_attr = paddle_op.attrs().get("scale")
|
||||
|
||||
trt_major = get_trt_version_list()[0]
|
||||
trt_minor = get_trt_version_list()[1]
|
||||
trt_version_float = float(f"{trt_major}.{trt_minor}")
|
||||
|
||||
resize_layer = network.add_resize(input_tensor)
|
||||
set_layer_name(resize_layer, paddle_op)
|
||||
# Set resize mode to LINEAR unconditionally
|
||||
if trt_version_float >= 8.6:
|
||||
resize_layer.resize_mode = trt.InterpolationMode.LINEAR
|
||||
else:
|
||||
resize_layer.resize_mode = trt.ResizeMode.LINEAR
|
||||
|
||||
# Set coordinate transformation based on align_corners and align_mode
|
||||
if align_corners:
|
||||
resize_layer.coordinate_transformation = (
|
||||
trt.ResizeCoordinateTransformation.ALIGN_CORNERS
|
||||
)
|
||||
else:
|
||||
if align_mode == 0:
|
||||
resize_layer.coordinate_transformation = (
|
||||
trt.ResizeCoordinateTransformation.HALF_PIXEL
|
||||
)
|
||||
else: # align_mode == 1
|
||||
resize_layer.coordinate_transformation = (
|
||||
trt.ResizeCoordinateTransformation.ASYMMETRIC
|
||||
)
|
||||
|
||||
if data_format == "NCHW":
|
||||
h_axis = 2
|
||||
w_axis = 3
|
||||
elif data_format == "NHWC":
|
||||
h_axis = 1
|
||||
w_axis = 2
|
||||
|
||||
in_dim = input_tensor.shape
|
||||
|
||||
outsize_tensor = None
|
||||
if trt_version_float >= 8.2:
|
||||
if not pir.is_fake_value(paddle_op.operands()[1].source()):
|
||||
size_tensor_operand = paddle_op.operands()[1].source()
|
||||
if len(inputs) > 1 and inputs[1] is not None:
|
||||
outsize_tensor = inputs[1]
|
||||
elif not pir.is_fake_value(paddle_op.operands()[2].source()):
|
||||
size_tensor_operand = paddle_op.operands()[2].source()
|
||||
size_tensor = inputs[2]
|
||||
if size_tensor_operand.is_combine():
|
||||
size_tensors = []
|
||||
if not isinstance(size_tensor, list):
|
||||
size_tensors = [size_tensor]
|
||||
else:
|
||||
size_tensors = size_tensor
|
||||
if len(size_tensors) >= 2:
|
||||
# Extract the first two elements representing height and width
|
||||
outsize_h = size_tensors[0]
|
||||
outsize_w = size_tensors[1]
|
||||
outsize_tensor = network.add_concatenation(
|
||||
[outsize_h, outsize_w]
|
||||
)
|
||||
set_layer_name(outsize_tensor, paddle_op)
|
||||
outsize_tensor = outsize_tensor.get_output(0)
|
||||
else:
|
||||
size_tensor_shape = size_tensor_operand.source().shape
|
||||
if size_tensor_shape.size >= 2:
|
||||
outsize_h = network.add_slice(
|
||||
size_tensor, start=[0], shape=[1], stride=[1]
|
||||
)
|
||||
set_layer_name(outsize_h, paddle_op)
|
||||
outsize_h = outsize_h.get_output(0)
|
||||
outsize_w = network.add_slice(
|
||||
size_tensor, start=[1], shape=[1], stride=[1]
|
||||
)
|
||||
set_layer_name(outsize_w, paddle_op)
|
||||
outsize_w = outsize_w.get_output(0)
|
||||
outsize_tensor = network.add_concatenation(
|
||||
[outsize_h, outsize_w]
|
||||
)
|
||||
set_layer_name(outsize_tensor, paddle_op)
|
||||
outsize_tensor = outsize_tensor.get_output(0)
|
||||
use_scales = True
|
||||
if outsize_tensor is not None:
|
||||
use_scales = False
|
||||
if outsize_tensor is None and len(scale_attr) == 0:
|
||||
use_scales = False
|
||||
|
||||
if use_scales:
|
||||
scale_h = -1.0
|
||||
scale_w = -1.0
|
||||
|
||||
if scale_attr and len(scale_attr) > 1:
|
||||
scale_h = scale_attr[0]
|
||||
scale_w = scale_attr[1]
|
||||
elif scale_attr and len(scale_attr) == 1:
|
||||
scale_h = scale_w = scale_attr[0]
|
||||
|
||||
if scale_w > 0 and scale_h > 0:
|
||||
if in_dim[h_axis] > 0 and in_dim[w_axis] > 0:
|
||||
out_h = int(in_dim[h_axis] * scale_h)
|
||||
out_w = int(in_dim[w_axis] * scale_w)
|
||||
else:
|
||||
if out_h > 0 and out_w > 0 and not (scale_w > 0 and scale_h > 0):
|
||||
if in_dim[h_axis] > 0 and in_dim[w_axis] > 0:
|
||||
scale_h = float(out_h) / float(in_dim[h_axis])
|
||||
scale_w = float(out_w) / float(in_dim[w_axis])
|
||||
|
||||
scales = [1.0] * len(input_tensor.shape)
|
||||
if data_format == "NCHW":
|
||||
scales[2] = scale_h
|
||||
scales[3] = scale_w
|
||||
elif data_format == "NHWC":
|
||||
scales[1] = scale_h
|
||||
scales[2] = scale_w
|
||||
|
||||
resize_layer.scales = scales
|
||||
else:
|
||||
if outsize_tensor is not None:
|
||||
outsize_itensors = []
|
||||
batch_dim = get_shape_tensor_element(
|
||||
network,
|
||||
input_shape_tensor,
|
||||
0,
|
||||
name=[paddle_op.name(), "batch_dim"],
|
||||
)
|
||||
outsize_itensors.append(batch_dim)
|
||||
if data_format == "NCHW":
|
||||
channel_dim = get_shape_tensor_element(
|
||||
network,
|
||||
input_shape_tensor,
|
||||
1,
|
||||
name=[paddle_op.name(), "channel_dim"],
|
||||
)
|
||||
outsize_itensors.append(channel_dim)
|
||||
outsize_itensors.append(outsize_tensor)
|
||||
elif data_format == "NHWC":
|
||||
channel_dim = get_shape_tensor_element(
|
||||
network,
|
||||
input_shape_tensor,
|
||||
3,
|
||||
name=[paddle_op.name(), "channel_dim"],
|
||||
)
|
||||
outsize_itensors.append(outsize_tensor)
|
||||
outsize_itensors.append(channel_dim)
|
||||
output_size_tensor = network.add_concatenation(outsize_itensors)
|
||||
set_layer_name(output_size_tensor, paddle_op)
|
||||
output_size_tensor = output_size_tensor.get_output(0)
|
||||
resize_layer.set_input(1, output_size_tensor)
|
||||
else:
|
||||
if data_format == "NCHW":
|
||||
shape_layer = network.add_shape(input_tensor)
|
||||
shape_output = shape_layer.get_output(0)
|
||||
# Get N and C from slice_layer output
|
||||
slice_layer = network.add_slice(
|
||||
shape_output, start=[0], shape=[2], stride=[1]
|
||||
)
|
||||
# Create H and W
|
||||
hw_constant = network.add_constant(
|
||||
shape=(2,),
|
||||
weights=trt.Weights(
|
||||
np.array([out_h, out_w], dtype=np.int32)
|
||||
),
|
||||
).get_output(0)
|
||||
# Create output shape(NCHW)
|
||||
concat_layer = network.add_concatenation(
|
||||
[slice_layer.get_output(0), hw_constant]
|
||||
)
|
||||
concat_layer.axis = 0
|
||||
resize_layer.set_input(1, concat_layer.get_output(0))
|
||||
elif data_format == "NHWC":
|
||||
shape_layer = network.add_shape(input_tensor)
|
||||
shape_output = shape_layer.get_output(0)
|
||||
# Get N and C from slice_layer output
|
||||
n_layer = network.add_slice(
|
||||
shape_output, start=[0], shape=[1], stride=[1]
|
||||
)
|
||||
c_layer = network.add_slice(
|
||||
shape_output, start=[3], shape=[1], stride=[1]
|
||||
)
|
||||
# Create H and W
|
||||
hw_constant = network.add_constant(
|
||||
shape=(2,),
|
||||
weights=trt.Weights(
|
||||
np.array([out_h, out_w], dtype=np.int32)
|
||||
),
|
||||
).get_output(0)
|
||||
# Create output shape(NHWC)
|
||||
concat_layer = network.add_concatenation(
|
||||
[n_layer.get_output(0), hw_constant, c_layer.get_output(0)]
|
||||
)
|
||||
concat_layer.axis = 0
|
||||
resize_layer.set_input(1, concat_layer.get_output(0))
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"Converter for bilinear_interp not support data_format {}.",
|
||||
data_format,
|
||||
)
|
||||
return resize_layer.get_output(0)
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.embedding")
|
||||
def embedding_converter(network, paddle_op, inputs):
|
||||
x = inputs[0]
|
||||
weight = inputs[1]
|
||||
gather_layer = network.add_gather(weight, x, 0)
|
||||
set_layer_name(gather_layer, paddle_op)
|
||||
return gather_layer.get_output(0)
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.unbind")
|
||||
def unbind_converter(network, paddle_op, inputs):
|
||||
x = inputs[0]
|
||||
input_shape = x.shape
|
||||
axis = paddle_op.attrs().get("axis")
|
||||
rank = len(input_shape)
|
||||
if axis < 0:
|
||||
axis += rank
|
||||
axis = int(axis)
|
||||
# Input for the add_slice layer
|
||||
start_tensors = []
|
||||
size_tensors = []
|
||||
# Input for the add_shuffle layer
|
||||
new_shape_tensors = []
|
||||
for i in range(rank):
|
||||
if axis == i:
|
||||
size_tensors.append(
|
||||
add_1D_constant_layer(
|
||||
network, 1, name=[paddle_op.name(), "size_tensor"]
|
||||
)
|
||||
)
|
||||
else:
|
||||
size_tensors.append(
|
||||
get_shape_tensor_element(
|
||||
network,
|
||||
trt_shape(network, x, name=[paddle_op.name(), "trt_shape"]),
|
||||
i,
|
||||
name=[paddle_op.name(), f"size_tensor_{i}"],
|
||||
)
|
||||
)
|
||||
new_shape_tensors.append(
|
||||
get_shape_tensor_element(
|
||||
network,
|
||||
trt_shape(network, x, name=[paddle_op.name(), "trt_shape"]),
|
||||
i,
|
||||
name=[paddle_op.name(), f"new_shape_tensor_{i}"],
|
||||
)
|
||||
)
|
||||
start_tensors.append(
|
||||
add_1D_constant_layer(
|
||||
network, 0, name=[paddle_op.name(), "start_tensor"]
|
||||
)
|
||||
)
|
||||
|
||||
new_shape_tensor = trt_concat(
|
||||
network, new_shape_tensors, name=[paddle_op.name(), "new_shape_tensor"]
|
||||
)
|
||||
stride = trt.Dims([1] * rank)
|
||||
outputs = []
|
||||
output_size = len(paddle_op.results()[0].type().as_vec_type().as_list())
|
||||
for i in range(output_size):
|
||||
start_tensors[axis] = add_1D_constant_layer(
|
||||
network, i, name=[paddle_op.name(), f"start_{i}_tensor"]
|
||||
)
|
||||
# Create Slice layer
|
||||
slice_layer = network.add_slice(
|
||||
x,
|
||||
stride,
|
||||
stride,
|
||||
stride,
|
||||
)
|
||||
slice_layer.set_input(1, trt_concat(network, start_tensors))
|
||||
slice_layer.set_input(2, trt_concat(network, size_tensors))
|
||||
set_layer_name(slice_layer, paddle_op)
|
||||
shuffle_layer = trt_reshape(
|
||||
network,
|
||||
slice_layer.get_output(0),
|
||||
new_shape_tensor,
|
||||
is_shape_tensor=True,
|
||||
name=[paddle_op.name(), f"shuffle_tensor_{i}"],
|
||||
)
|
||||
outputs.append(shuffle_layer)
|
||||
return outputs
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.nearest_interp")
|
||||
def nearest_interp_converter(network, paddle_op, inputs):
|
||||
input_tensor = inputs[0]
|
||||
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)
|
||||
input_rank = input_shape_tensor.shape # The reason is unknown that adding this unused code make input_shape_tensor maintain the correct result.
|
||||
data_format = paddle_op.attrs().get("data_format")
|
||||
interp_method = paddle_op.attrs().get("interp_method")
|
||||
align_corners = paddle_op.attrs().get("align_corners")
|
||||
out_h = paddle_op.attrs().get("out_h")
|
||||
out_w = paddle_op.attrs().get("out_w")
|
||||
out_d = paddle_op.attrs().get("out_d")
|
||||
scale_attr = paddle_op.attrs().get("scale")
|
||||
|
||||
# Parse TensorRT version
|
||||
trt_major = get_trt_version_list()[0]
|
||||
trt_minor = get_trt_version_list()[1]
|
||||
trt_version_float = float(f"{trt_major}.{trt_minor}")
|
||||
|
||||
# Create Resize layer
|
||||
resize_layer = network.add_resize(input_tensor)
|
||||
set_layer_name(resize_layer, paddle_op)
|
||||
|
||||
if trt_version_float >= 8.6:
|
||||
if align_corners:
|
||||
resize_layer.coordinate_transformation = (
|
||||
trt.ResizeCoordinateTransformation.ASYMMETRIC
|
||||
)
|
||||
else:
|
||||
resize_layer.coordinate_transformation = (
|
||||
trt.ResizeCoordinateTransformation.ASYMMETRIC
|
||||
)
|
||||
|
||||
in_dim = input_tensor.shape
|
||||
scale_h = 1.0
|
||||
scale_w = 1.0
|
||||
|
||||
if scale_attr is not None and len(scale_attr) >= 2:
|
||||
scale_h = scale_attr[0]
|
||||
scale_w = scale_attr[1]
|
||||
else:
|
||||
if out_h > 0 and out_w > 0:
|
||||
if data_format == "NCHW":
|
||||
h_axis = 2
|
||||
w_axis = 3
|
||||
elif data_format == "NHWC":
|
||||
h_axis = 1
|
||||
w_axis = 2
|
||||
|
||||
scale_h = float(out_h) / float(in_dim[h_axis])
|
||||
scale_w = float(out_w) / float(in_dim[w_axis])
|
||||
|
||||
outsize_tensor = None
|
||||
if inputs[2] is not None:
|
||||
outsize_tensor = network.add_concatenation(inputs[2])
|
||||
set_layer_name(outsize_tensor, paddle_op)
|
||||
outsize_tensor = outsize_tensor.get_output(0)
|
||||
|
||||
scales = [1.0] * len(input_tensor.shape)
|
||||
if data_format == "NCHW":
|
||||
scales[1] = 1.0
|
||||
scales[2] = scale_h
|
||||
scales[3] = scale_w
|
||||
elif data_format == "NHWC":
|
||||
scales[1] = scale_h
|
||||
scales[2] = scale_w
|
||||
scales[3] = 1.0
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported data format {data_format}, only NCHW or NHWC are supported."
|
||||
)
|
||||
if outsize_tensor is not None:
|
||||
outsize_itensors = []
|
||||
batch_dim = get_shape_tensor_element(
|
||||
network, input_shape_tensor, 0, name=[paddle_op.name(), "batch_dim"]
|
||||
)
|
||||
outsize_itensors.append(batch_dim)
|
||||
if data_format == "NCHW":
|
||||
channel_dim = get_shape_tensor_element(
|
||||
network,
|
||||
input_shape_tensor,
|
||||
1,
|
||||
name=[paddle_op.name(), "channel_dim"],
|
||||
)
|
||||
outsize_itensors.append(channel_dim)
|
||||
outsize_itensors.append(outsize_tensor)
|
||||
elif data_format == "NHWC":
|
||||
channel_dim = get_shape_tensor_element(
|
||||
network,
|
||||
input_shape_tensor,
|
||||
3,
|
||||
name=[paddle_op.name(), "channel_dim"],
|
||||
)
|
||||
outsize_itensors.append(outsize_tensor)
|
||||
outsize_itensors.append(channel_dim)
|
||||
resize_layer.set_input(
|
||||
1, network.add_concatenation(outsize_itensors).get_output(0)
|
||||
)
|
||||
else:
|
||||
resize_layer.scales = scales
|
||||
|
||||
return resize_layer.get_output(0)
|
||||
|
||||
|
||||
@converter_registry.register("pd_op.linear_interp")
|
||||
def linear_interp_converter(network, paddle_op, inputs):
|
||||
input_tensor = inputs[0]
|
||||
data_layout = paddle_op.attrs().get("data_format")
|
||||
interp_method = paddle_op.attrs().get("interp_method")
|
||||
align_corners = paddle_op.attrs().get("align_corners")
|
||||
out_w = paddle_op.attrs().get("out_w")
|
||||
scale_attr = paddle_op.attrs().get("scale")
|
||||
layer = network.add_resize(input_tensor)
|
||||
set_layer_name(layer, paddle_op)
|
||||
trt_major = get_trt_version_list()[0]
|
||||
trt_minor = get_trt_version_list()[1]
|
||||
trt_version_float = float(f"{trt_major}.{trt_minor}")
|
||||
|
||||
if trt_version_float >= 8.6:
|
||||
layer.resize_mode = trt.InterpolationMode.LINEAR
|
||||
else:
|
||||
layer.resize_mode = trt.ResizeMode.LINEAR
|
||||
|
||||
if align_corners:
|
||||
layer.coordinate_transformation = (
|
||||
trt.ResizeCoordinateTransformation.ALIGN_CORNERS
|
||||
)
|
||||
else:
|
||||
layer.coordinate_transformation = (
|
||||
trt.ResizeCoordinateTransformation.HALF_PIXEL
|
||||
)
|
||||
|
||||
in_dim = input_tensor.shape
|
||||
scale_w = -1.0
|
||||
|
||||
if scale_attr and len(scale_attr) > 0:
|
||||
scale_w = scale_attr[0]
|
||||
|
||||
w_axis = 2 if data_layout == "NCHW" else 1
|
||||
|
||||
if float(scale_w) > 0.0:
|
||||
out_w = int(in_dim[w_axis] * scale_w)
|
||||
|
||||
outsize_tensor = None
|
||||
if len(inputs) > 1 and inputs[1] is not None:
|
||||
outsize_tensor = inputs[1]
|
||||
|
||||
if outsize_tensor is None:
|
||||
if len(inputs) > 2 and inputs[2] is not None:
|
||||
outsize_tensor = inputs[2][0]
|
||||
|
||||
if out_w > 0 and scale_w <= 0:
|
||||
scale_w = float(out_w) / float(in_dim[w_axis])
|
||||
|
||||
scales = [1.0]
|
||||
if data_layout == "NCHW":
|
||||
scales.append(1.0)
|
||||
scales.append(scale_w)
|
||||
elif data_layout == "NHWC":
|
||||
scales.append(scale_w)
|
||||
scales.append(1.0)
|
||||
|
||||
if outsize_tensor is not None:
|
||||
outsize_itensors = []
|
||||
input_shape = trt_shape(
|
||||
network, input_tensor, name=[paddle_op.name(), "input_shape"]
|
||||
)
|
||||
batch_dim = get_shape_tensor_element(
|
||||
network, input_shape, 0, name=[paddle_op.name(), "batch_dim"]
|
||||
)
|
||||
outsize_itensors.append(batch_dim)
|
||||
|
||||
if data_layout == "NCHW":
|
||||
channel_dim = get_shape_tensor_element(
|
||||
network, input_shape, 1, name=[paddle_op.name(), "channel_dim"]
|
||||
)
|
||||
outsize_itensors.append(channel_dim)
|
||||
outsize_itensors.append(outsize_tensor)
|
||||
elif data_layout == "NHWC":
|
||||
outsize_itensors.append(outsize_tensor)
|
||||
channel_dim = get_shape_tensor_element(
|
||||
network, input_shape, 2, name=[paddle_op.name(), "channel_dim"]
|
||||
)
|
||||
outsize_itensors.append(channel_dim)
|
||||
|
||||
layer.set_input(
|
||||
1,
|
||||
trt_concat(
|
||||
network,
|
||||
outsize_itensors,
|
||||
name=[paddle_op.name(), "outsize_itensors"],
|
||||
),
|
||||
)
|
||||
else:
|
||||
layer.scales = scales
|
||||
|
||||
return layer.get_output(0)
|
||||
Reference in New Issue
Block a user