wehub-resource-sync
557 lines
20 KiB
Python
557 lines
20 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 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)
|