wehub-resource-sync
1220 lines
41 KiB
Python
1220 lines
41 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 inspect
|
|
import logging
|
|
import os
|
|
import sys
|
|
|
|
import numpy as np
|
|
import tensorrt as trt
|
|
|
|
from paddle.tensorrt.util import TensorRTConfigManager, TensorRTConstantManager
|
|
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
|
|
if parent_dir not in sys.path:
|
|
sys.path.append(parent_dir)
|
|
|
|
|
|
from tensorrt import INetworkDefinition, ITensor
|
|
|
|
from paddle.base.log_helper import get_logger
|
|
|
|
_logger = get_logger(
|
|
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
|
)
|
|
from paddle.base.libpaddle.pir import (
|
|
get_attrs_map_json,
|
|
get_inputs_type_json,
|
|
get_outputs_type_json,
|
|
)
|
|
|
|
version = trt.__version__
|
|
version_list = list(map(int, version.split('.')))
|
|
|
|
|
|
def has_dynamic_shape(shape):
|
|
return any(s == -1 for s in shape)
|
|
|
|
|
|
def append_ones(network, input, name, num_prepend_ones):
|
|
layer = network.add_shuffle(input)
|
|
|
|
if has_dynamic_shape(input.shape):
|
|
input_shape_layer = network.add_shape(input)
|
|
prepend_shape_layer = network.add_constant(
|
|
(num_prepend_ones,), np.ones((num_prepend_ones,), dtype=np.int32)
|
|
)
|
|
reshape_dim_layer = network.add_concatenation(
|
|
[prepend_shape_layer.get_output(0), input_shape_layer.get_output(0)]
|
|
)
|
|
reshape_dim_layer.axis = 0
|
|
layer.set_input(1, reshape_dim_layer.get_output(0))
|
|
if name is not None:
|
|
set_layer_name(input_shape_layer, [name[0], "input_shape_layer"])
|
|
set_layer_name(
|
|
prepend_shape_layer, [name[0], "prepend_shape_layer"]
|
|
)
|
|
set_layer_name(reshape_dim_layer, [name[0], "reshape_dim_layer"])
|
|
else:
|
|
layer.reshape_dims = (1,) * num_prepend_ones + tuple(input.shape)
|
|
|
|
if name is not None:
|
|
set_layer_name(layer, name)
|
|
return layer.get_output(0)
|
|
|
|
|
|
def broadcast(network, a, b, a_name, b_name, paddle_op, preset_diff=0):
|
|
a_shape = tuple(a.shape)
|
|
b_shape = tuple(b.shape)
|
|
|
|
diff = len(a_shape) - len(b_shape) - preset_diff
|
|
if diff > 0:
|
|
b = append_ones(network, b, [paddle_op.name(), b_name], diff)
|
|
elif diff < 0:
|
|
a = append_ones(network, a, [paddle_op.name(), a_name], -diff)
|
|
|
|
return a, b
|
|
|
|
|
|
def get_axes_for_reduce_op(
|
|
dim,
|
|
has_implicit_batch_dimension=False,
|
|
):
|
|
if isinstance(dim, int):
|
|
dim = (dim,)
|
|
|
|
if has_implicit_batch_dimension:
|
|
assert 0 not in dim, (
|
|
"Can't reduce over batch dimension when it's implicit."
|
|
)
|
|
|
|
axes = 0
|
|
for d in dim:
|
|
axes |= 1 << (d - (1 if has_implicit_batch_dimension else 0))
|
|
|
|
return axes
|
|
|
|
|
|
def get_dynamic_dims(shape):
|
|
"""
|
|
This function finds the dynamic dimensions in the given
|
|
shape. A dimension is dynamic if it's -1.
|
|
|
|
Args:
|
|
shape (Shape): A sequence of integer that represents
|
|
the shape of a tensor.
|
|
|
|
Returns:
|
|
A list of integers contains all the dynamic dimensions
|
|
in the given shape
|
|
"""
|
|
dynamic_dims = []
|
|
for i, s in enumerate(shape):
|
|
if s == -1:
|
|
dynamic_dims.append(i)
|
|
return dynamic_dims
|
|
|
|
|
|
def get_trt_plugin(plugin_name, field_collection, version, plugin_namespace=""):
|
|
plugin_registry = trt.get_plugin_registry()
|
|
plugin_creator = plugin_registry.get_plugin_creator(
|
|
plugin_name, version, plugin_namespace
|
|
)
|
|
assert plugin_creator, (
|
|
f"Unable to found plugin creator with name {plugin_name}"
|
|
)
|
|
plugin = plugin_creator.create_plugin(
|
|
name=plugin_name, field_collection=field_collection
|
|
)
|
|
assert plugin is not None, f"Plugin:{plugin_name} could not be fetched"
|
|
return plugin
|
|
|
|
|
|
def get_positive_dim(dim, dim_size):
|
|
if dim < 0:
|
|
return dim % dim_size
|
|
return dim
|
|
|
|
|
|
def add_elementwise_layer(network, paddle_op, inputs, op_type):
|
|
from paddle.tensorrt.util import support_fp32_mix_precision
|
|
|
|
weight_shape = paddle_op.operands()[1].source().shape
|
|
input_shape = paddle_op.operands()[0].source().shape
|
|
|
|
weight_tensor = inputs[1]
|
|
input_tensor = inputs[0]
|
|
if type(inputs[1]) == trt.Weights:
|
|
weight_tensor = network.add_constant(weight_shape, inputs[1])
|
|
set_layer_name(weight_tensor, paddle_op)
|
|
weight_tensor = weight_tensor.get_output(0)
|
|
if type(inputs[0]) == trt.Weights:
|
|
input_tensor = network.add_constant(input_shape, inputs[0])
|
|
set_layer_name(input_tensor, paddle_op)
|
|
input_tensor = input_tensor.get_output(0)
|
|
lhs_val, rhs_val = broadcast(
|
|
network,
|
|
input_tensor,
|
|
weight_tensor,
|
|
"input_tensor_broadcast",
|
|
"weight_tensor_broadcast",
|
|
paddle_op,
|
|
)
|
|
layer = network.add_elementwise(lhs_val, rhs_val, op_type)
|
|
set_layer_name(layer, paddle_op)
|
|
support_fp32_mix_precision(paddle_op.name(), layer)
|
|
return layer.get_output(0)
|
|
|
|
|
|
# Create and add 1D constant layer
|
|
def add_1D_constant_layer(
|
|
network, data, dtype=np.int32, is_scalar=False, name=None
|
|
):
|
|
if not isinstance(data, list):
|
|
data = [data]
|
|
shape = () if is_scalar else (len(data),)
|
|
constant_data = np.array(data, dtype=dtype)
|
|
constant_layer = network.add_constant(shape, constant_data)
|
|
set_layer_name(constant_layer, name)
|
|
return constant_layer.get_output(0)
|
|
|
|
|
|
# Create and add ND constant layer
|
|
def add_constant_layer(network, data, shape, dtype=np.int32, name=None):
|
|
constant_data = np.array(data, dtype=dtype)
|
|
constant_data = np.resize(constant_data, shape)
|
|
constant_layer = network.add_constant(shape, constant_data)
|
|
set_layer_name(constant_layer, name)
|
|
return constant_layer.get_output(0)
|
|
|
|
|
|
# Create an constant layer with shape_tensor and value
|
|
def fill_constant_layer(
|
|
network, shape_tensor, tensor_rank, data, trt_dtype, name=None
|
|
):
|
|
fill_layer = network.add_fill(
|
|
trt.Dims([tensor_rank]), trt.FillOperation.LINSPACE
|
|
)
|
|
np_dtype = map_trt_dtype(trt_dtype)
|
|
fill_layer.set_input(0, shape_tensor)
|
|
fill_layer.set_input(
|
|
1, add_1D_constant_layer(network, data, np_dtype, is_scalar=True)
|
|
)
|
|
beta = [0] * tensor_rank
|
|
fill_layer.set_input(
|
|
2, add_1D_constant_layer(network, beta, np_dtype, is_scalar=False)
|
|
)
|
|
set_layer_name(fill_layer, name)
|
|
return fill_layer.get_output(0)
|
|
|
|
|
|
def trt_expand(network, input, rank, shape_tensor, shape_rank, name=None):
|
|
def process_names(name, layer_name):
|
|
if name is not None:
|
|
return [name[0], layer_name]
|
|
else:
|
|
return None
|
|
|
|
if rank < shape_rank:
|
|
one_rank_tensor = add_1D_constant_layer(
|
|
network,
|
|
[1] * (shape_rank - rank),
|
|
name=process_names(name, "one_rank_tensor"),
|
|
)
|
|
in_shape_tensor = trt_shape(
|
|
network, input, name=process_names(name, "in_shape_tensor")
|
|
)
|
|
itensors = [one_rank_tensor, in_shape_tensor]
|
|
input_shape_tensor = trt_concat(
|
|
network, itensors, name=process_names(name, "input_shape_tensor")
|
|
)
|
|
else:
|
|
input_shape_tensor = trt_shape(
|
|
network, input, name=process_names(name, "input_shape_tensor")
|
|
)
|
|
|
|
new_input_tensor = trt_reshape(
|
|
network,
|
|
input,
|
|
input_shape_tensor,
|
|
process_names(name, "new_input_tensor"),
|
|
True,
|
|
)
|
|
|
|
start = [0] * shape_rank
|
|
starts_tensor = add_1D_constant_layer(
|
|
network, start, name=process_names(name, "starts_tensor")
|
|
)
|
|
one_tensor = add_1D_constant_layer(
|
|
network, 1, name=process_names(name, "one_tensor")
|
|
)
|
|
sizes_tensor = trt_max(
|
|
network,
|
|
input_shape_tensor,
|
|
shape_tensor,
|
|
name=process_names(name, "sizes_tensor"),
|
|
)
|
|
input_sub_tensor = trt_sub(
|
|
network,
|
|
input_shape_tensor,
|
|
one_tensor,
|
|
name=process_names(name, "input_sub_tensor"),
|
|
)
|
|
strides_tensor = trt_min(
|
|
network,
|
|
one_tensor,
|
|
input_sub_tensor,
|
|
name=process_names(name, "strides_tensor"),
|
|
)
|
|
|
|
slice_layer = network.add_slice(
|
|
new_input_tensor, start, [0] * len(start), [0] * len(start)
|
|
)
|
|
slice_layer.set_input(1, starts_tensor)
|
|
slice_layer.set_input(2, sizes_tensor)
|
|
slice_layer.set_input(3, strides_tensor)
|
|
set_layer_name(slice_layer, name)
|
|
|
|
return slice_layer.get_output(0)
|
|
|
|
|
|
# Concat not make rank changed
|
|
def trt_concat(network, inputs, axis=0, name=None):
|
|
concat_layer = network.add_concatenation(inputs=inputs)
|
|
if axis != 0:
|
|
concat_layer.axis = axis
|
|
set_layer_name(concat_layer, name)
|
|
return concat_layer.get_output(0)
|
|
|
|
|
|
def trt_cast(network, input, dtype, name=None):
|
|
identity_layer = network.add_identity(input)
|
|
identity_layer.set_output_type(0, dtype)
|
|
identity_layer.get_output(0).dtype = dtype
|
|
set_layer_name(identity_layer, name)
|
|
return identity_layer.get_output(0)
|
|
|
|
|
|
def trt_shape(
|
|
network: INetworkDefinition, input: ITensor, name=None
|
|
) -> ITensor:
|
|
"""
|
|
Add a IShapeLayer to get the shape of `input` ITensor.
|
|
This includes a workaround that casting the shape result(int64) from TRT10 back to int32.
|
|
Many existing paddle op kernels only support input shape tensor as int32
|
|
, to make TRT op more compatible with other paddle op, we cast back to int32.
|
|
NOTE: please remove this workaround when all paddle op supports shape tensor in int64
|
|
"""
|
|
shape_layer = network.add_shape(input)
|
|
set_layer_name(shape_layer, name)
|
|
if version_list[0] >= 10: # trt_version >=10
|
|
# workaround
|
|
if name is not None:
|
|
name = [name[0], "trt_cast"]
|
|
return trt_cast(
|
|
network, shape_layer.get_output(0), trt.int32, name=name
|
|
)
|
|
return shape_layer.get_output(0)
|
|
|
|
|
|
def trt_reshape(network, input, new_shape, name=None, is_shape_tensor=False):
|
|
reshape_layer = network.add_shuffle(input)
|
|
if is_shape_tensor:
|
|
reshape_layer.set_input(1, new_shape)
|
|
else:
|
|
reshape_layer.reshape_dims = new_shape
|
|
if name is not None:
|
|
if isinstance(name, list):
|
|
set_layer_name(reshape_layer, name)
|
|
else:
|
|
reshape_layer.name = name
|
|
return reshape_layer.get_output(0)
|
|
|
|
|
|
# resize shape tensor's shape to 1dim
|
|
def resize_to_1d(network, shape_tensor, name=None):
|
|
if shape_tensor is None:
|
|
return shape_tensor
|
|
if len(shape_tensor.shape) > 1:
|
|
# shape_tensor need 1-dim in trt
|
|
shape_tensor_layer = network.add_shuffle(shape_tensor)
|
|
numel = 1
|
|
for ele in shape_tensor.shape:
|
|
numel *= ele
|
|
shape_tensor_layer.reshape_dims = [numel]
|
|
set_layer_name(shape_tensor_layer, name)
|
|
shape_tensor = shape_tensor_layer.get_output(0)
|
|
return shape_tensor
|
|
|
|
|
|
# Get element tensor of 1D shape tensor
|
|
def get_shape_tensor_element(network, x, index, is_scalar=False, name=None):
|
|
assert index >= 0, (
|
|
f"The index should be greater or equal than 0, but got {index}"
|
|
)
|
|
index_tensor_name = [name[0], "index_tensor"] if name is not None else None
|
|
index_tensor = add_1D_constant_layer(
|
|
network, index, is_scalar=is_scalar, name=index_tensor_name
|
|
)
|
|
gather_layer = network.add_gather(input=x, indices=index_tensor, axis=0)
|
|
if name is not None:
|
|
set_layer_name(gather_layer, [name[0], "gather_layer"])
|
|
shape_tensor = resize_to_1d(network, gather_layer.get_output(0), name=name)
|
|
return shape_tensor
|
|
|
|
|
|
def trt_less(network, a, b, name=None):
|
|
layer = network.add_elementwise(a, b, trt.ElementWiseOperation.LESS)
|
|
set_layer_name(layer, name)
|
|
return layer.get_output(0)
|
|
|
|
|
|
def trt_sum(network, a, b, name=None):
|
|
layer = network.add_elementwise(a, b, trt.ElementWiseOperation.SUM)
|
|
set_layer_name(layer, name)
|
|
return layer.get_output(0)
|
|
|
|
|
|
def trt_max(network, a, b, name=None):
|
|
layer = network.add_elementwise(a, b, trt.ElementWiseOperation.MAX)
|
|
set_layer_name(layer, name)
|
|
return layer.get_output(0)
|
|
|
|
|
|
def trt_sub(network, a, b, name=None):
|
|
layer = network.add_elementwise(a, b, trt.ElementWiseOperation.SUB)
|
|
set_layer_name(layer, name)
|
|
return layer.get_output(0)
|
|
|
|
|
|
def trt_min(network, a, b, name=None):
|
|
layer = network.add_elementwise(a, b, trt.ElementWiseOperation.MIN)
|
|
set_layer_name(layer, name)
|
|
return layer.get_output(0)
|
|
|
|
|
|
def trt_div(network, a, b, name=None):
|
|
layer = network.add_elementwise(a, b, trt.ElementWiseOperation.DIV)
|
|
set_layer_name(layer, name)
|
|
return layer.get_output(0)
|
|
|
|
|
|
def trt_floor_div(network, a, b, name=None):
|
|
layer = network.add_elementwise(a, b, trt.ElementWiseOperation.FLOOR_DIV)
|
|
set_layer_name(layer, name)
|
|
return layer.get_output(0)
|
|
|
|
|
|
def trt_equal(network, a, b, name=None):
|
|
layer = network.add_elementwise(a, b, trt.ElementWiseOperation.EQUAL)
|
|
set_layer_name(layer, name)
|
|
return layer.get_output(0)
|
|
|
|
|
|
def trt_gather(network, input, indices, axis=0, name=None):
|
|
if name is not None:
|
|
name = [name[0], "indices_tensor"]
|
|
indices_tensor = add_1D_constant_layer(network, indices, name=name)
|
|
gather_layer = network.add_gather(input, indices_tensor, axis)
|
|
set_layer_name(gather_layer, name)
|
|
result = gather_layer.get_output(0)
|
|
return result
|
|
|
|
|
|
def trt_prod(network, a, b, name=None):
|
|
layer = network.add_elementwise(a, b, trt.ElementWiseOperation.PROD)
|
|
set_layer_name(layer, name)
|
|
return layer.get_output(0)
|
|
|
|
|
|
def trt_pow(network, a, b, name=None):
|
|
layer = network.add_elementwise(a, b, trt.ElementWiseOperation.POW)
|
|
set_layer_name(layer, name)
|
|
return layer.get_output(0)
|
|
|
|
|
|
def cast_tensor(network, input_tensor, dtype, name=None):
|
|
layer = network.add_identity(input_tensor)
|
|
layer.set_output_type(0, dtype)
|
|
set_layer_name(layer, name)
|
|
return layer.get_output(0)
|
|
|
|
|
|
def build_start_tensor(network, rank, axis_tensor, offset, name=None):
|
|
# Create indices_tensor [0, 1, ..., rank-1]
|
|
indices = np.arange(rank, dtype=np.int32)
|
|
indices_name = [name[0], "indices_tensor"] if name is not None else None
|
|
indices_tensor = network.add_constant([rank], indices)
|
|
set_layer_name(indices_tensor, indices_name)
|
|
indices_tensor = indices_tensor.get_output(0)
|
|
|
|
# Create mask: mask = (indices == axis_tensor)
|
|
mask_name = [name[0], "mask"] if name is not None else None
|
|
mask = network.add_elementwise(
|
|
indices_tensor, axis_tensor, trt.ElementWiseOperation.EQUAL
|
|
)
|
|
set_layer_name(mask, mask_name)
|
|
mask = mask.get_output(0)
|
|
mask_int = cast_tensor(
|
|
network,
|
|
mask,
|
|
trt.int32,
|
|
name=[name[0], "mask_int"] if name is not None else None,
|
|
)
|
|
|
|
# Calculate start_tensor = mask_int * offset
|
|
start_tensor = network.add_elementwise(
|
|
mask_int, offset, trt.ElementWiseOperation.PROD
|
|
)
|
|
set_layer_name(start_tensor, name)
|
|
start_tensor = start_tensor.get_output(0)
|
|
|
|
return start_tensor
|
|
|
|
|
|
def build_size_tensor(
|
|
network, rank, axis_tensor, size_value, input_shape_tensor, name=None
|
|
):
|
|
# Create indices_tensor [0, 1, ..., rank-1]
|
|
indices = np.arange(rank, dtype=np.int32)
|
|
indices_name = [name[0], 'indices_tensor'] if name is not None else None
|
|
indices_tensor = network.add_constant([rank], indices)
|
|
set_layer_name(indices_tensor, indices_name)
|
|
indices_tensor = indices_tensor.get_output(0)
|
|
|
|
# Create mask: mask = (indices == axis_tensor)
|
|
mask_name = [name[0], 'mask'] if name is not None else None
|
|
mask = network.add_elementwise(
|
|
indices_tensor, axis_tensor, trt.ElementWiseOperation.EQUAL
|
|
)
|
|
set_layer_name(mask, mask_name)
|
|
mask = mask.get_output(0)
|
|
mask_int = cast_tensor(
|
|
network,
|
|
mask,
|
|
trt.int32,
|
|
name=[name[0], "mask_int"] if name is not None else None,
|
|
)
|
|
|
|
# Create ones_tensor
|
|
ones_name = [name[0], 'ones_tensor'] if name is not None else None
|
|
ones_tensor = network.add_constant([rank], np.ones([rank], dtype=np.int32))
|
|
set_layer_name(ones_tensor, ones_name)
|
|
ones_tensor = ones_tensor.get_output(0)
|
|
|
|
# Calculate inverse_mask = ones_tensor - mask_int
|
|
inverse_mask_name = [name[0], 'inverse_mask'] if name is not None else None
|
|
inverse_mask = network.add_elementwise(
|
|
ones_tensor, mask_int, trt.ElementWiseOperation.SUB
|
|
)
|
|
set_layer_name(inverse_mask, inverse_mask_name)
|
|
inverse_mask = inverse_mask.get_output(0)
|
|
|
|
# Calculate size_tensor = mask_int * size_value + inverse_mask * input_shape_tensor
|
|
size_value_broadcast_name = (
|
|
[name[0], 'size_value_broadcast'] if name is not None else None
|
|
)
|
|
size_value_broadcast = network.add_elementwise(
|
|
mask_int, size_value, trt.ElementWiseOperation.PROD
|
|
)
|
|
set_layer_name(size_value_broadcast, size_value_broadcast_name)
|
|
size_value_broadcast = size_value_broadcast.get_output(0)
|
|
|
|
input_shape_broadcast_name = (
|
|
[name[0], 'input_shape_broadcast'] if name is not None else None
|
|
)
|
|
input_shape_broadcast = network.add_elementwise(
|
|
inverse_mask, input_shape_tensor, trt.ElementWiseOperation.PROD
|
|
)
|
|
set_layer_name(input_shape_broadcast, input_shape_broadcast_name)
|
|
input_shape_broadcast = input_shape_broadcast.get_output(0)
|
|
|
|
size_tensor = network.add_elementwise(
|
|
size_value_broadcast,
|
|
input_shape_broadcast,
|
|
trt.ElementWiseOperation.SUM,
|
|
)
|
|
set_layer_name(size_tensor, name)
|
|
size_tensor = size_tensor.get_output(0)
|
|
|
|
return size_tensor
|
|
|
|
|
|
# convert trt_dtype to numpy dtype
|
|
def map_trt_dtype(trt_dtype):
|
|
dtype_map = {
|
|
trt.DataType.FLOAT: np.float32,
|
|
trt.DataType.HALF: np.float16,
|
|
trt.DataType.INT32: np.int32,
|
|
trt.DataType.INT8: np.int8,
|
|
trt.DataType.BOOL: bool,
|
|
}
|
|
if trt_dtype in dtype_map:
|
|
return dtype_map[trt_dtype]
|
|
else:
|
|
raise TypeError(f"Unsupported trt_dtype: {trt_dtype}")
|
|
|
|
|
|
# Reduce the given tensor in the TensorRT network to a scalar
|
|
def trt_reduce_to_scalar(network, tensor, dtype=trt.int32, name=None):
|
|
if len(tensor.shape) == 0:
|
|
return tensor
|
|
axes = 0
|
|
for i in range(len(tensor.shape)):
|
|
axes |= 1 << i
|
|
reduce_layer = network.add_reduce(
|
|
tensor, trt.ReduceOperation.SUM, axes, keep_dims=False
|
|
)
|
|
if name is not None:
|
|
set_layer_name(reduce_layer, [name[0], 'reduce_layer'])
|
|
scalar_name = name
|
|
scalar = trt_cast(
|
|
network, reduce_layer.get_output(0), dtype, name=scalar_name
|
|
)
|
|
return scalar
|
|
|
|
|
|
def convert_conv2d(network, paddle_op, inputs):
|
|
from paddle.tensorrt.util import (
|
|
RefitManager,
|
|
RefitRole,
|
|
support_fp32_mix_precision,
|
|
)
|
|
|
|
bias = None
|
|
if (
|
|
paddle_op.name() == "pd_op.conv2d"
|
|
or paddle_op.name() == "pd_op.depthwise_conv2d"
|
|
):
|
|
input_tensor, filter = inputs
|
|
elif (
|
|
paddle_op.name() == "pd_op.conv2d_transpose"
|
|
or paddle_op.name() == "pd_op.depthwise_conv2d_transpose"
|
|
):
|
|
if len(inputs) == 3:
|
|
input_tensor, filter, output_size = inputs
|
|
elif len(inputs) == 2:
|
|
input_tensor, filter = inputs
|
|
output_size = None
|
|
else:
|
|
raise ValueError("Invalid number of inputs for conv2d_transpose")
|
|
if paddle_op.name() == "pd_op.fused_conv2d_add_act":
|
|
input_tensor, filter, bias, _ = inputs
|
|
input_shape = paddle_op.operands()[0].source().shape
|
|
filter_shape = paddle_op.operands()[1].source().shape
|
|
|
|
if len(filter_shape) != 4:
|
|
raise ValueError(
|
|
f"filter's dims size should be 4, but got {len(filter_shape)}"
|
|
)
|
|
|
|
n_output = filter_shape[0]
|
|
n_input = filter_shape[1]
|
|
filter_h = filter_shape[2]
|
|
filter_w = filter_shape[3]
|
|
|
|
paddings = paddle_op.attrs().get("paddings", [0, 0])
|
|
stride = paddle_op.attrs().get("strides", [1, 1])
|
|
dilation = paddle_op.attrs().get("dilations", [1, 1])
|
|
groups = paddle_op.attrs().get("groups", 1)
|
|
|
|
if has_dynamic_shape(input_shape):
|
|
assert input_shape[1] != -1, (
|
|
"Channel dim can't be dynamic for transpose convolution."
|
|
)
|
|
|
|
output_padding = paddle_op.attrs().get("output_padding", [0, 0])
|
|
padding_algorithm = paddle_op.attrs().get("padding_algorithm", "EXPLICIT")
|
|
if padding_algorithm == "VALID":
|
|
paddings = [0] * len(paddings)
|
|
|
|
nv_ksize = trt.DimsHW(filter_h, filter_w)
|
|
nv_dilations = trt.DimsHW(dilation[0], dilation[1])
|
|
nv_strides = trt.DimsHW(stride[0], stride[1])
|
|
|
|
pre_paddings = [0, 0]
|
|
post_paddings = [0, 0]
|
|
|
|
if isinstance(filter, trt.Weights):
|
|
weight_filter = filter
|
|
else:
|
|
weight_filter = trt.Weights()
|
|
|
|
if len(paddings) == 2:
|
|
pre_paddings[0] = paddings[0]
|
|
pre_paddings[1] = paddings[1]
|
|
post_paddings[0] = paddings[0]
|
|
post_paddings[1] = paddings[1]
|
|
elif len(paddings) == 4:
|
|
pre_paddings[0] = paddings[0]
|
|
pre_paddings[1] = paddings[2]
|
|
post_paddings[0] = paddings[1]
|
|
post_paddings[1] = paddings[3]
|
|
else:
|
|
raise ValueError(f"Unsupported paddings size: {len(paddings)}")
|
|
|
|
if paddle_op.name() == "pd_op.fused_conv2d_add_act":
|
|
constant_manager = TensorRTConstantManager()
|
|
bias_source_op = paddle_op.operands()[2].source().get_defining_op()
|
|
if bias_source_op.name() == "builtin.parameter":
|
|
bias_name = bias_source_op.attrs()['parameter_name']
|
|
elif bias_source_op.name() == "builtin.constant":
|
|
bias_np = bias_source_op.attrs()['value']
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported bias source op: {bias_source_op.name()}"
|
|
)
|
|
bias_np = constant_manager.get_constant_value(bias_name)
|
|
bias_weights = trt.Weights(bias_np)
|
|
layer = network.add_convolution_nd(
|
|
input=input_tensor,
|
|
num_output_maps=n_output,
|
|
kernel_shape=nv_ksize,
|
|
kernel=weight_filter,
|
|
bias=bias_weights,
|
|
)
|
|
elif (
|
|
paddle_op.name() == "pd_op.conv2d"
|
|
or paddle_op.name() == "pd_op.depthwise_conv2d"
|
|
):
|
|
layer = network.add_convolution_nd(
|
|
input=input_tensor,
|
|
num_output_maps=n_output,
|
|
kernel_shape=nv_ksize,
|
|
kernel=weight_filter,
|
|
bias=None,
|
|
)
|
|
elif (
|
|
paddle_op.name() == "pd_op.conv2d_transpose"
|
|
or paddle_op.name() == "pd_op.depthwise_conv2d_transpose"
|
|
):
|
|
layer = network.add_deconvolution_nd(
|
|
input=input_tensor,
|
|
num_output_maps=n_input * groups,
|
|
kernel_shape=nv_ksize,
|
|
kernel=weight_filter,
|
|
bias=None,
|
|
)
|
|
|
|
if isinstance(filter, trt.ITensor):
|
|
layer.set_input(1, filter)
|
|
layer.stride_nd = nv_strides
|
|
layer.pre_padding = pre_paddings
|
|
|
|
if output_padding:
|
|
post_paddings[0] -= output_padding[0]
|
|
post_paddings[1] -= output_padding[1]
|
|
|
|
if post_paddings[0] < 0 or post_paddings[1] < 0:
|
|
raise ValueError("The value PostPadding should be >= 0.")
|
|
|
|
layer.post_padding = post_paddings
|
|
layer.num_groups = groups
|
|
|
|
if padding_algorithm == "SAME":
|
|
layer.padding_mode = trt.PaddingMode.SAME_UPPER
|
|
nv_dilations = trt.DimsHW(1, 1)
|
|
|
|
layer.dilation_nd = nv_dilations
|
|
set_layer_name(layer, paddle_op)
|
|
support_fp32_mix_precision(paddle_op.name(), layer)
|
|
|
|
trt_manager = TensorRTConfigManager()
|
|
if trt_manager.get_refit_params_path():
|
|
filter_param = paddle_op.operands()[1].source()
|
|
filter_name = filter_param.get_defining_op().attrs()['parameter_name']
|
|
refit_manager = RefitManager()
|
|
refit_manager.set_mapping(filter_name, filter_name, RefitRole.CONSTANT)
|
|
|
|
return layer.get_output(0)
|
|
|
|
|
|
def convert_conv3d(network, paddle_op, inputs):
|
|
from paddle.tensorrt.util import (
|
|
RefitManager,
|
|
RefitRole,
|
|
)
|
|
|
|
input_tensor, filter = inputs
|
|
filter_shape = paddle_op.operands()[1].source().shape
|
|
|
|
n_output = filter_shape[0]
|
|
n_input = filter_shape[1]
|
|
filter_d = filter_shape[2]
|
|
filter_h = filter_shape[3]
|
|
filter_w = filter_shape[4]
|
|
|
|
if isinstance(filter, trt.Weights):
|
|
weight_filter = filter
|
|
else:
|
|
weight_filter = trt.Weights()
|
|
|
|
groups = paddle_op.attrs().get("groups", 1)
|
|
dilations = paddle_op.attrs().get("dilations", [1, 1, 1])
|
|
strides = paddle_op.attrs().get("strides", [1, 1, 1])
|
|
paddings = paddle_op.attrs().get("paddings", [0, 0, 0])
|
|
padding_algorithm = paddle_op.attrs().get("padding_algorithm", "EXPLICIT")
|
|
# for conv3d_transpose
|
|
output_padding = paddle_op.attrs().get("output_padding", [])
|
|
|
|
nv_ksize = trt.Dims3(filter_d, filter_h, filter_w)
|
|
nv_dilations = trt.Dims3(dilations[0], dilations[1], dilations[2])
|
|
nv_strides = trt.Dims3(strides[0], strides[1], strides[2])
|
|
nv_pre_paddings = trt.Dims3(paddings[0], paddings[1], paddings[2])
|
|
|
|
if paddle_op.name() == "pd_op.conv3d":
|
|
layer = network.add_convolution_nd(
|
|
input=input_tensor,
|
|
num_output_maps=n_output,
|
|
kernel_shape=nv_ksize,
|
|
kernel=weight_filter,
|
|
bias=None,
|
|
)
|
|
elif paddle_op.name() == "pd_op.conv3d_transpose":
|
|
layer = network.add_deconvolution_nd(
|
|
input=input_tensor,
|
|
num_output_maps=n_input * groups,
|
|
kernel_shape=nv_ksize,
|
|
kernel=weight_filter,
|
|
bias=None,
|
|
)
|
|
layer.stride_nd = nv_strides
|
|
layer.pre_padding = nv_pre_paddings
|
|
|
|
nv_post_paddings = trt.Dims3(paddings[0], paddings[1], paddings[2])
|
|
if output_padding:
|
|
nv_post_paddings[0] -= output_padding[0]
|
|
nv_post_paddings[1] -= output_padding[1]
|
|
nv_post_paddings[2] -= output_padding[2]
|
|
|
|
if (
|
|
nv_post_paddings[0] < 0
|
|
or nv_post_paddings[1] < 0
|
|
or nv_post_paddings[2] < 0
|
|
):
|
|
raise ValueError(
|
|
"The value in conv3d_transpose's PostPadding should be >= 0."
|
|
)
|
|
if isinstance(filter, trt.ITensor):
|
|
layer.set_input(1, filter)
|
|
|
|
layer.post_padding = nv_post_paddings
|
|
layer.num_groups = groups
|
|
|
|
if padding_algorithm == "SAME":
|
|
layer.padding_mode = trt.PaddingMode.SAME_UPPER
|
|
|
|
layer.dilation_nd = nv_dilations
|
|
set_layer_name(layer, paddle_op)
|
|
trt_manager = TensorRTConfigManager()
|
|
if trt_manager.get_refit_params_path():
|
|
filter_param = paddle_op.operands()[1].source()
|
|
filter_name = filter_param.get_defining_op().attrs()['parameter_name']
|
|
refit_manager = RefitManager()
|
|
refit_manager.set_mapping(filter_name, filter_name, RefitRole.CONSTANT)
|
|
|
|
return layer.get_output(0)
|
|
|
|
|
|
def get_input_constant_value(paddle_op, inputs, input_index):
|
|
input_op = paddle_op.operands()[input_index].source().get_defining_op()
|
|
constant_manager = TensorRTConstantManager()
|
|
if input_op.name() == "builtin.constant":
|
|
return constant_manager.get_constant_value(
|
|
input_op.attrs()["value"]
|
|
).tolist()
|
|
elif input_op.name() == "pd_op.full_int_array":
|
|
return input_op.attrs()["value"]
|
|
elif input_op.name() == "pd_op.full":
|
|
return [input_op.attrs()["value"]]
|
|
else:
|
|
return None
|
|
|
|
|
|
def add_reduce_layer(network, paddle_op, inputs, op_type):
|
|
input_tensor = inputs[0]
|
|
axis = get_input_constant_value(paddle_op, inputs, 1)
|
|
input_shape = paddle_op.operands()[0].source().shape
|
|
keepdim = paddle_op.attrs()["keepdim"]
|
|
if network.has_implicit_batch_dimension:
|
|
assert axis != 0, (
|
|
"can't reduce on axis == 0 when network has implicit batch dimension"
|
|
)
|
|
output_shape = []
|
|
if len(axis) == 0:
|
|
axis = list(range(len(input_shape)))
|
|
for i in range(len(axis)):
|
|
if axis[i] < 0:
|
|
axis[i] = len(input_shape) + axis[i]
|
|
layer = network.add_reduce(
|
|
input_tensor,
|
|
op_type,
|
|
axes=get_axes_for_reduce_op(axis),
|
|
keep_dims=keepdim,
|
|
)
|
|
set_layer_name(layer, paddle_op)
|
|
layer.get_output(0).dtype = layer.get_input(0).dtype
|
|
return layer.get_output(0)
|
|
|
|
|
|
def add_cast_reduce_layer(network, paddle_op, inputs, op_type):
|
|
input_tensor = inputs[0]
|
|
cast_layer = network.add_identity(input_tensor)
|
|
set_layer_name(cast_layer, paddle_op)
|
|
cast_layer.set_output_type(0, trt.int32)
|
|
cast_layer.get_output(0).dtype = trt.int32
|
|
|
|
axis = paddle_op.attrs().get("axis")
|
|
input_shape = paddle_op.operands()[0].source().shape
|
|
input_dims = len(input_shape)
|
|
keepdim = paddle_op.attrs().get("keepdim")
|
|
|
|
if len(axis) == 0:
|
|
axes = 0
|
|
for i in range(input_dims):
|
|
axes |= 1 << i
|
|
else:
|
|
for i in range(len(axis)):
|
|
if axis[i] < 0:
|
|
axis[i] += input_dims
|
|
|
|
axes = get_axes_for_reduce_op(axis)
|
|
|
|
reduce_layer = network.add_reduce(
|
|
cast_layer.get_output(0),
|
|
op_type,
|
|
axes=axes,
|
|
keep_dims=keepdim,
|
|
)
|
|
set_layer_name(reduce_layer, paddle_op)
|
|
bool_layer = network.add_identity(reduce_layer.get_output(0))
|
|
set_layer_name(bool_layer, paddle_op)
|
|
bool_layer.set_output_type(0, trt.bool)
|
|
bool_layer.get_output(0).dtype = trt.bool
|
|
return bool_layer.get_output(0)
|
|
|
|
|
|
def fix_negative_indices(network, input_shape, indices, name=None):
|
|
rank = len(input_shape.shape)
|
|
zero_tensor_name = [name[0], 'zero_tensor'] if name else None
|
|
zero_tensor = add_1D_constant_layer(
|
|
network, [0] * rank, name=zero_tensor_name
|
|
)
|
|
minus_one_tensor_name = [name[0], 'minus_one_tensor'] if name else None
|
|
minus_one_tensor = add_1D_constant_layer(
|
|
network, [-1] * rank, name=minus_one_tensor_name
|
|
)
|
|
|
|
min_indices_zero_name = [name[0], 'min_indices_zero'] if name else None
|
|
min_indices_zero = trt_min(
|
|
network, indices, zero_tensor, name=min_indices_zero_name
|
|
)
|
|
sign_name = [name[0], 'sign'] if name else None
|
|
sign = trt_max(network, min_indices_zero, minus_one_tensor, name=sign_name)
|
|
sub_name = [name[0], 'sub'] if name else None
|
|
sub = trt_prod(network, sign, input_shape, name=sub_name)
|
|
fixed_indices = trt_sub(network, indices, sub, name=name)
|
|
return fixed_indices
|
|
|
|
|
|
def trt_unsqueeze(network, input_tensor, axes, name=None):
|
|
input_shape_name = [name[0], 'input_shape'] if name else None
|
|
input_shape = network.add_shape(input_tensor)
|
|
set_layer_name(input_shape, input_shape_name)
|
|
input_shape = input_shape.get_output(0)
|
|
|
|
axis_set = set(axes)
|
|
|
|
subscripts = list(range(len(input_tensor.shape)))
|
|
|
|
for axis in sorted(axis_set):
|
|
subscripts.insert(axis, len(input_tensor.shape))
|
|
|
|
one_tensor_name = [name[0], 'one_tensor'] if name else None
|
|
one_tensor = network.add_constant((1,), np.array([1], dtype=np.int32))
|
|
set_layer_name(one_tensor, one_tensor_name)
|
|
one_tensor = one_tensor.get_output(0)
|
|
extended_shape_name = [name[0], 'extended_shape'] if name else None
|
|
extended_shape = network.add_concatenation(
|
|
[input_shape, one_tensor],
|
|
)
|
|
set_layer_name(extended_shape, extended_shape_name)
|
|
extended_shape = extended_shape.get_output(0)
|
|
|
|
gather_layer_name = [name[0], 'gather_layer'] if name else None
|
|
gather_layer = network.add_gather(
|
|
extended_shape,
|
|
network.add_constant(
|
|
(len(subscripts),), np.array(subscripts, dtype=np.int32)
|
|
).get_output(0),
|
|
axis=0,
|
|
)
|
|
set_layer_name(gather_layer, gather_layer_name)
|
|
new_shape_tensor = gather_layer.get_output(0)
|
|
|
|
reshaped_tensor = network.add_shuffle(input_tensor)
|
|
reshaped_tensor.set_input(1, new_shape_tensor)
|
|
set_layer_name(reshaped_tensor, name)
|
|
|
|
return reshaped_tensor.get_output(0)
|
|
|
|
|
|
def squeeze_trt(network, input_tensor, axes, name=None):
|
|
input_shape_name = [name[0], 'input_shape'] if name else None
|
|
input_shape = network.add_shape(input_tensor)
|
|
set_layer_name(input_shape, input_shape_name)
|
|
input_shape = input_shape.get_output(0)
|
|
input_shape = input_tensor.shape
|
|
all_dims = list(range(len(input_shape)))
|
|
remaining_dims = [dim for dim in all_dims if dim not in axes]
|
|
|
|
input_shape_tensor_name = [name[0], 'input_shape_tensor'] if name else None
|
|
input_shape_tensor = network.add_shape(input_tensor)
|
|
set_layer_name(input_shape_tensor, input_shape_tensor_name)
|
|
input_shape_tensor = input_shape_tensor.get_output(0)
|
|
|
|
remaining_dims_tensor_name = (
|
|
[name[0], 'remaining_dims_tensor'] if name else None
|
|
)
|
|
remaining_dims_tensor = network.add_constant(
|
|
(len(remaining_dims),), np.array(remaining_dims, dtype=np.int32)
|
|
)
|
|
set_layer_name(remaining_dims_tensor, remaining_dims_tensor_name)
|
|
remaining_dims_tensor = remaining_dims_tensor.get_output(0)
|
|
|
|
new_shape_tensor_name = [name[0], 'new_shape_tensor'] if name else None
|
|
new_shape_tensor = network.add_gather(
|
|
input_shape_tensor, remaining_dims_tensor, axis=0
|
|
)
|
|
set_layer_name(new_shape_tensor, new_shape_tensor_name)
|
|
new_shape_tensor = new_shape_tensor.get_output(0)
|
|
reshape_layer = network.add_shuffle(input_tensor)
|
|
reshape_layer.set_input(1, new_shape_tensor)
|
|
set_layer_name(reshape_layer, name)
|
|
return reshape_layer.get_output(0)
|
|
|
|
|
|
def unary_op_converter(network, paddle_op, inputs):
|
|
from paddle.tensorrt import PrecisionMode
|
|
|
|
ops_type_map = {
|
|
"pd_op.sqrt": [trt.UnaryOperation.SQRT],
|
|
"pd_op.sqrt_": [trt.UnaryOperation.SQRT],
|
|
"pd_op.floor": [trt.UnaryOperation.FLOOR],
|
|
"pd_op.exp": [trt.UnaryOperation.EXP],
|
|
"pd_op.abs": [trt.UnaryOperation.ABS],
|
|
"pd_op.abs_": [trt.UnaryOperation.ABS],
|
|
"pd_op.sin": [trt.UnaryOperation.SIN],
|
|
"pd_op.cos": [trt.UnaryOperation.COS],
|
|
"pd_op.sinh": [trt.UnaryOperation.SINH],
|
|
"pd_op.cosh": [trt.UnaryOperation.COSH],
|
|
"pd_op.asinh": [trt.UnaryOperation.ASINH],
|
|
"pd_op.acosh": [trt.UnaryOperation.ACOSH],
|
|
"pd_op.atanh": [trt.UnaryOperation.ATANH],
|
|
"pd_op.ceil": [trt.UnaryOperation.CEIL],
|
|
"pd_op.reciprocal": [trt.UnaryOperation.RECIP],
|
|
"pd_op.erf": [trt.UnaryOperation.ERF],
|
|
"pd_op.sign": [trt.UnaryOperation.SIGN],
|
|
"pd_op.round": [trt.UnaryOperation.ROUND],
|
|
"pd_op.logical_not": [trt.UnaryOperation.NOT],
|
|
"pd_op.rsqrt": [trt.UnaryOperation.SQRT, trt.UnaryOperation.RECIP],
|
|
"pd_op.tan": [trt.UnaryOperation.TAN],
|
|
"pd_op.asin": [trt.UnaryOperation.ASIN],
|
|
"pd_op.acos": [trt.UnaryOperation.ACOS],
|
|
"pd_op.atan": [trt.UnaryOperation.ATAN],
|
|
}
|
|
|
|
input_tensor = inputs[0]
|
|
layer = None
|
|
org_type = input_tensor.dtype
|
|
|
|
trt_type_mapping = {
|
|
trt.DataType.INT8: trt.int8,
|
|
trt.DataType.INT32: trt.int32,
|
|
}
|
|
|
|
trt_manager = TensorRTConfigManager()
|
|
precision_mode = trt_manager.get_precision_mode()
|
|
|
|
need_cast = org_type in [trt.DataType.INT8, trt.DataType.INT32]
|
|
if need_cast:
|
|
identity_layer = network.add_identity(input_tensor)
|
|
if precision_mode == PrecisionMode.FP32:
|
|
identity_layer.set_output_type(0, trt.float32)
|
|
else:
|
|
identity_layer.set_output_type(0, trt.float16)
|
|
set_layer_name(identity_layer, paddle_op)
|
|
input_tensor = identity_layer.get_output(0)
|
|
|
|
if paddle_op.name() in ops_type_map:
|
|
for trt_op in ops_type_map[paddle_op.name()]:
|
|
layer = network.add_unary(input_tensor, trt_op)
|
|
set_layer_name(layer, paddle_op)
|
|
input_tensor = layer.get_output(0)
|
|
else:
|
|
raise NotImplementedError(
|
|
f"Unsupported unary operation: {paddle_op.name()}"
|
|
)
|
|
if need_cast:
|
|
restore_layer = network.add_identity(input_tensor)
|
|
restore_layer.set_output_type(0, trt_type_mapping[org_type])
|
|
set_layer_name(restore_layer, paddle_op)
|
|
input_tensor = restore_layer.get_output(0)
|
|
|
|
return input_tensor
|
|
|
|
|
|
# get the length of the specified axis for input_tensor
|
|
def get_axis_length(network, input_tensor, axis, is_scalar=False, name=None):
|
|
input_shape = input_tensor.shape
|
|
if input_shape[axis] >= 0:
|
|
output_tensor = add_1D_constant_layer(
|
|
network, input_shape[axis], is_scalar=is_scalar, name=name
|
|
)
|
|
else:
|
|
shape_name = [name[0], 'dynamic_shape'] if name else None
|
|
dynamic_shape = trt_shape(network, input_tensor, name=shape_name)
|
|
output_tensor = get_shape_tensor_element(
|
|
network, dynamic_shape, axis, is_scalar, name=name
|
|
)
|
|
return output_tensor
|
|
|
|
|
|
def WithFp16():
|
|
from paddle.tensorrt import PrecisionMode
|
|
|
|
trt_manager = TensorRTConfigManager()
|
|
precision_mode = trt_manager.get_precision_mode()
|
|
enable_fp16 = False
|
|
if precision_mode == PrecisionMode.FP16:
|
|
enable_fp16 = True
|
|
# TODO(lizexu123) WithInt8() and use_dla are not yet implemented
|
|
return enable_fp16
|
|
|
|
|
|
def set_layer_name(layer, second_param):
|
|
"""
|
|
Sets standardized names for converter output layers following the format: `<id>_<pd_op>-><layerName>(<inputIds>)`
|
|
|
|
Naming Rule:
|
|
Format: <sequence_number>_<paddle_op_name>-><layer_variable_name>(<comma_separated_input_ids>)
|
|
Components:
|
|
- sequence_number: Output tensor's unique ID from layer
|
|
- paddle_op_name: Name of source Paddle operator
|
|
- layer_variable_name: Variable name referencing the layer in code
|
|
- input_ids: Input tensor IDs from preceding layers
|
|
|
|
Args:
|
|
layer (ILayer): Target layer to name
|
|
second_param: Context-dependent parameter:
|
|
- For non-public functions: paddle_op (op object)
|
|
- For public functions: [paddle_op_name (str), layer_var_name (str)] list
|
|
- When name=None in public functions: Enables nested handling
|
|
"""
|
|
if second_param is not None:
|
|
if isinstance(second_param, list):
|
|
# Handling for public function layer
|
|
op_name, layer_var_name = second_param
|
|
else:
|
|
# Handling for layer
|
|
op_name = second_param.name()
|
|
layer_var_name = None
|
|
if op_name is not None:
|
|
# Retrieve the name of the variable that refers to the layer
|
|
for (
|
|
var_name,
|
|
var_val,
|
|
) in inspect.currentframe().f_back.f_locals.items():
|
|
if var_val is layer:
|
|
layer_var_name = var_name
|
|
break
|
|
|
|
# Retrieve the input id of the layer
|
|
if op_name is not None and layer_var_name is not None:
|
|
input_ids = []
|
|
i = 0
|
|
while (input_tensor := layer.get_input(i)) is not None:
|
|
input_name = input_tensor.name
|
|
if "Unnamed Layer" in input_name:
|
|
input_id = input_name.split("*")[1].split(")")[0].strip()
|
|
else:
|
|
input_id = input_name
|
|
input_ids.append(input_id)
|
|
i += 1
|
|
|
|
# Retrieve the output id of the layer
|
|
output_name = layer.get_output(0).name
|
|
if "Unnamed Layer" in output_name:
|
|
sequence_number = (
|
|
output_name.split("*")[1].split(")")[0].strip()
|
|
)
|
|
else:
|
|
sequence_number = output_name
|
|
|
|
formatted_name = (
|
|
f"{sequence_number}_"
|
|
f"{op_name}->"
|
|
f"{layer_var_name}"
|
|
f"({', '.join(input_ids)})"
|
|
)
|
|
layer.name = formatted_name
|
|
|
|
|
|
def generic_plugin_converter(network, paddle_op, inputs, extra_attrs=None):
|
|
op_name = paddle_op.name()
|
|
|
|
if extra_attrs is not None:
|
|
attrs_map_info = get_attrs_map_json(extra_attrs)
|
|
else:
|
|
attrs_map_info = get_attrs_map_json(paddle_op)
|
|
|
|
input_type_info = get_inputs_type_json(paddle_op)
|
|
output_type_info = get_outputs_type_json(paddle_op)
|
|
|
|
plugin_fields = [
|
|
trt.PluginField(
|
|
"op_name",
|
|
np.array(list(op_name), dtype=np.bytes_),
|
|
trt.PluginFieldType.CHAR,
|
|
),
|
|
trt.PluginField(
|
|
"attrs_map_info",
|
|
np.array(list(attrs_map_info), dtype=np.bytes_),
|
|
trt.PluginFieldType.CHAR,
|
|
),
|
|
trt.PluginField(
|
|
"inputs_type_info",
|
|
np.array(list(input_type_info), dtype=np.bytes_),
|
|
trt.PluginFieldType.CHAR,
|
|
),
|
|
trt.PluginField(
|
|
"outputs_type_info",
|
|
np.array(list(output_type_info), dtype=np.bytes_),
|
|
trt.PluginFieldType.CHAR,
|
|
),
|
|
]
|
|
|
|
plugin_field_collection = trt.PluginFieldCollection(plugin_fields)
|
|
|
|
plugin_name = "pir_generic_plugin"
|
|
plugin_version = "1"
|
|
plugin = get_trt_plugin(
|
|
plugin_name, plugin_field_collection, plugin_version
|
|
)
|
|
|
|
layer = network.add_plugin_v2(inputs, plugin)
|
|
return layer
|