o + i @sddlmZddlZddlZddlmZddlZddlZddl m Z m Z ddl m Z ddlmZddlmZddlmZmZmZdd lmZdd lmZdd lmZdd lmZmZmZm Z m!Z!ee"ej#d dZ$GdddZ%GdddeZ&GdddZ'ddZ(dddZ)ddZ*dS)) annotationsN)Enum)coredygraph) scope_guard)Variable) get_logger)_get_function_names_from_layerget_ast_static_function to_static)StaticFunction)Layer)PaddleToTensorRTConverter)forbid_op_lower_trtmark_builtin_op run_pir_passrun_trt_partitionwarmup_shape_inferz&%(asctime)s-%(levelname)s: %(message)s)fmtc@s.eZdZ       ddddZddZdS)InputNfloat32 warmup_datatuple[np.ndarray, ...] | Nonemin_input_shape tuple | Nonemax_input_shapeoptim_input_shapeinput_data_type str | None input_rangenamereturnNonecCs||dur|s |s |rtd|dus|durtdn d|||fvr'td||_||_||_||_||_||_||_ dS)a A class used to configure input data for models. This class serves two purposes: 1. Random Data Generation: When no input data is supplied, it automatically generates random input data based on the specified minimum, optimal, and maximum shapes. In this mode,you can configure the data type (e.g., 'float32', 'int64', etc.) and the range of values (e.g.,(0.0, 1.0) for floats or (1, 10) for integers). 2. User-Provided Input: Alternatively, you can supply your own input data via the `warmup_data` argument. In this case, the provided data will be used directly, and the`input_data_type` and `input_range` settings will be ignored. Args: warmup_data (tuple): The tuple of actual input data (for the automatic shape collection mechanism). min_input_shape (tuple): The shape of the minimum input tensor. max_input_shape (tuple): The shape of the maximum input tensor. optim_input_shape (tuple): The shape of the optimal input tensor. input_data_type (str, optional): The data type for the input tensors, such as 'float32' or 'int64' or 'float32' or 'int32' (default is float32). This option only applies when min_input_shape, optim_input_shape, and max_input_shape are provided; it does not apply to warmup_data. input_range (tuple, optional): The range of values used to generate input data. For floats, the default range is (0.0, 1.0). For integers, the default range is (1, 10). This option only applies when min_input_shape, optim_input_shape, and max_input_shape are provided; it does not apply to warmup_data. name:(str,optional): The name of the input to the model. Returns: None Examples: .. code-block:: python >>> # example 1: >>> from paddle.tensorrt.export import Input >>> input_config = Input( >>> min_input_shape=(1,100), >>> optim_input_shape=(4,100), >>> max_input_shape=(8,100), >>> ) >>> input_config.input_data_type='int64' >>> input_config.input_range=(1,10) >>> # example 2: >>> from paddle.tensorrt.export import Input >>> import numpy as np >>> input_config = Input( >>> warmup_data=( >>> np.random.rand(1,100).astype(np.float32), >>> np.random.rand(4,100).astype(np.float32), >>> np.random.rand(8,100).astype(np.float32), >>> ) >>> ) Nz7warmup data provided; min/max/optim shapes are ignored.zWhen warmup_data is provided,input_data_type and input_range are ignored.These parameters only apply whtn generate random data using min/opt/max shapes.z?When warm_data is None, min/max/optim shapes must be specified.) ValueError_loggerwarningrrrrrrr )selfrrrrrrr r'b/home/app/PaddleOCR-VL-test/.venv_paddleocr/lib/python3.10/site-packages/paddle/tensorrt/export.py__init__6s*=  zInput.__init__cCsd|jdur td|jdurd|jvrdnd|_|j\}}||krJtj|j||jd|_tj|j||jd|_ tj|j ||jd|_ |j|j |j fSd|jvr}tj j |||jd|j|_tj j |||jd|j|_ tj j |||j d|j|_ n-tj j|||jd|j|_tj j|||jd|j|_ tj j|||j d|j|_ |j|j |j fS) aF Generates random input data based on the user-specified min_input_shape, optim_input_shape, and max_input_shape, as well as the data type and input range. Returns: tuple(numpy.ndarray, numpy.ndarray, numpy.ndarray): A tuple containing the generated input data for the minimum, optimal, and maximum shapes. Examples: .. code-block:: python >>> from paddle.tensorrt.export import Input >>> input_config = Input( >>> min_input_shape=(1,100), >>> optim_input_shape=(4,100), >>> max_input_shape=(8,100), >>> ) >>> input.input_data_type='int64' >>> input.input_range=(1,10) >>> input_min_data, input_optim_data, input_max_data = input_config.generate_input_data() NzHgenerate_input_data() should not be called when warmup_data is provided.float)gg?) )Zdtypeint)size)r RuntimeErrorrrnpfullrZinput_min_datarZinput_optim_datarZinput_max_datarandomrandintZastypeuniform)r&lowhighr'r'r(generate_input_datasr       zInput.generate_input_data)NNNNrNN)rrrrrrrrrrrrr rr!r")__name__ __module__ __qualname__r)r7r'r'r'r(r5s Urc@seZdZdZdZdZdZdS) PrecisionModeFP32FP16BF16INT8N)r8r9r:r<r=r>r?r'r'r'r(r;s r;c @s0eZdZdddejddgddddf dddZdS)TensorRTConfigNi@FTinputslistmin_subgraph_size int | Nonesave_model_dirr disable_opsstr | list | Noneprecision_moder; ops_run_floatoptimization_leveldisable_passesworkspace_sizeuse_cuda_graph bool | Nonerefit_params_pathdisable_logglingr!r"c Csdd|D} t| rt| std||_||_||_||_||_||_||_ ||_ | |_ | |_ | |_ | |_|j r@|j dtjd|idS)aj A class for configuring TensorRT optimizations. Args: inputs (list): A list of Input configurations min_subgraph_size (int, optional): The minimum number of operations in a subgraph for TensorRT to optimize (default is 3). save_model_dir (str, optional): The directory where the optimized model will be saved (default is not to save). disable_ops : (str|list, optional): A string representing the names of operations that should not be entering by TensorRT (default is None). precision_mode (PrecisionMode, optional): Specifies the precision mode for TensorRT optimization. The options are: - PrecisionMode.FP32: 32-bit floating point precision (default). - PrecisionMode.FP16: 16-bit floating point precision. - PrecisionMode.INT8: 8-bit integer precision. - PrecisionMode.BFP16: 16-bit Brain Floating Point precision. Only supported in TensorRT versions greater than 9.0. ops_run_float (str|list, optional): A set of operation names that should be executed using FP32 precision regardless of the `tensorrt_precision_mode` setting. optimization_level (int, optional): Set TensorRT optimization level (default is 3). Only supported in TensorRT versions greater than 8.6. disable_passes : (str|list, optional): A list of string representing the names of pass that should not be used for origin program (default is []). workspace_size (int, optional): Specifies the maximum GPU memory (in bytes) that TensorRT can use for the optimization process (default is 1 << 30). use_cuda_graph (bool, optional): Specify whether TensorRT enables cuda_graph during the optimization process (default is false). refit_params_path(str, optional): The path to the weights that need to be refitted. disable_loggling (bool, optional): Specifies whether to enable GLOG info output during the optimization process (default is true). Returns: None Examples: .. code-block:: python >>> # example 1: >>> from paddle.tensorrt.export import ( >>> Input, >>> TensorRTConfig, >>> PrecisionMode, >>> ) >>> input_config = Input( >>> min_input_shape=(1,100), >>> optim_input_shape=(4,100), >>> max_input_shape=(8,100), >>> ) >>> input_config.input_data_type='int64' >>> input_config.input_range=(1,10) >>> trt_config = TensorRTConfig(inputs=[input_config]) >>> trt_config.disable_ops = ["pd_op.dropout"] >>> trt_config.precision_mode = PrecisionMode.FP16 >>> trt_config.ops_run_float = "pd_op.conv2d" >>> trt_config.workspace_size = 1 << 32 >>> # example 2: >>> from paddle.tensorrt.export import ( >>> Input, >>> TensorRTConfig, >>> PrecisionMode, >>> ) >>> input_config = Input( >>> warmup_data=( >>> np.random.rand(1,100).astype(np.float32), >>> np.random.rand(4,100).astype(np.float32), >>> np.random.rand(8,100).astype(np.float32), >>> ) >>> ) >>> trt_config = TensorRTConfig(inputs=[input_config]) cSsg|]}|jduqS)Nr.0ir'r'r( @sz+TensorRTConfig.__init__..z,All Inputs must have input_data if any does.Zconstant_folding_passZFLAGS_trt_min_group_sizeN)anyallr#rBrDrFrIrJrGrLrKrMrNrPrQappendpaddleZ frameworkZ set_flags)r&rBrDrFrGrIrJrKrLrMrNrPrQZhas_input_datar'r'r(r)s*X zTensorRTConfig.__init__)rBrCrDrErFrrGrHrIr;rJrHrKrErLrCrMrErNrOrPrrQrOr!r")r8r9r:r;r<r)r'r'r'r(r@sr@c Cst|tjjjjstdt|g}|j D]}| dks'| dkr2| d}| |qtj "g}|jdjdurdd|jD}tdd |Dd ksYJt|d}t|D]'} i} t|jD]\} } | j durx| j n|| } || | | | <ql| | qcn8d d|jD}tt|dD]'} i} t|jD]\}} | j dur| j n||} ||| | | <q| | qt||j||jd }t|||d }|jrt||jt|t|}t|||d}|g}|j D]}| dkr|D] }|}| |qq|j rRg}|!dd|j Dt"d}tj#$|}t%|tj#j&|j ||||dWdn 1sMwY|WdS1s_wYdS)NzEprogram type must be paddle.base.libpaddle.pir.Program, but received pd_op.data pd_op.feedr rcSsg|]}|jqSr'rR)rTinpr'r'r(rVhsz"convert_to_trt..cSsh|]}t|qSr')len)rTtr'r'r( jz!convert_to_trt..r+cSsg|]}|qSr')r7rSr'r'r(rVsra)rLscoperI)feedsrb) trt_configz pd_op.fetchcss:|]}|dks|dkr|D]}|VqqdS)r[r\N)r results)rTopresultr'r'r( sz!convert_to_trt..)program)' isinstancerZbaseZ libpaddlepirZProgram TypeErrortypeZ global_blockopsr attrsrY pir_utilsIrGuardrBrr^range enumeraterrLrIrrGrrrrZconvert_program_to_trtZoperandssourcerFextend CUDAPlacestaticExecutorrZsave_inference_model)rirdrbZ feed_namerf param_namercZ input_tuplesZ num_samplesZ sample_idxZ feed_dictrUr]r jZprogram_with_pir converterZtrt_output_varoperandruZ input_valuesplaceexer'r'r(convert_to_trtYs             &rc Kstjj|fd|i|}tot}t|tjr|j}n|}d}|durt|t rKt |D]}t ||d} t| t rJd|krJt dt|dq/t|ttfs\tdt|dg}tj|D]+} t| tjjrs|| qdt| tjjttjjfr|tjj| qd|| qdd} t} i} t|t rttt |}t|}|j s|j!rd} n|g}g}d}|D]}t|t rt"t ||d} t| t r| j#r| }|||j$j%d|fq| j&|| |d }nZd|kr |rtj'||}t|j(|dd }|j&| |d }d}n:qt|t r/t"|} | j#r'| }||| fq| j&||d }nt"|} |r=tj'||}t| |dd }|j)}d}t|t rT|*}nt|t rc| j+rc| j+*}|r~i}i}|,D]\}} ||| j-<| || j-<qnt.//t0|j1D]!\}}|j2sq| 3|j-4}||j-54}|6|qWdn 1swYqtj78"|j9}|j:}tj;j||| }|| fWdS1swYdS)ai Convert a dynamic graph API to a static graph and apply TensorRT optimizations if relevant parameters are configured. Args: function (callable): Callable dynamic graph function. If it used as a decorator, the decorated function will be parsed as this parameter. input_spec (list[InputSpec]|tuple[InputSpec]): list/tuple of InputSpec to specific the shape/dtype/name information of each input Tensor. config: (TensorRTConfig): The configuration of TensorRTConfig. kwargs: Support keys including `property`, set `property` to True if the function is python property. Returns: tuple: A tuple containing two elements. The first element is the TensorRT optimized program., optionally optimized with TensorRT if configured. The second element is the scope containing the parameters. input_specNforwardzIf there are static functions other than 'forward' that need to be saved, the input 'input_spec' should be None, but received the type of 'input_spec' is .zIThe input input_spec should be 'list', but received input_spec's type is FT) with_hook is_prim_infer)rZ full_graph)rfetch)riZ fetch_listZfetch_var_name)?rZZjitr rZ_is_fwd_prim_enabledZ_is_bwd_prim_enabledrjZ DataParallelZ_layersr r getattrr r#rnrCtuplermutilsflattenrxZ InputSpecrYeagerZTensorrrlValueZ from_tensorZScopesetsortedZ_forward_pre_hooksZ_forward_post_hooksr Z is_property __class__r8Z#concrete_program_specify_input_specZpack_sequence_asrconcrete_programZto_static_state_dictZclass_instanceitemsr rguardzip parametersZ persistablevarZ get_tensorvalueZ_share_data_withrqrr main_programZoutputsrkexecutorZ_add_pir_fetch_opsr) functionrconfigkwargsZ static_netrZ inner_layerZinner_input_spec member_nameZ static_funcrrrbZextra_var_infoZ functionsZ property_valsrZ attr_funcZ immediate_valZstatic_forwardZstatic_functionZdygraph_state_dictZstate_names_dictZstate_var_dictZstrcutured_nameZtensorrZparam_or_buffer_tensorZ src_tensorrZ output_varsZprogram_with_trtr'r'r( _convert_s                     &rcCstj|jtj|krtdtj}td}tj |}d}tj |r[|}tj |\}}tj |\}} | dkrBd}n| dkrId}ntd| dtj ||d } n$|}|d } tj|drld}ntj|drwd}ntd |d tj| std | d |rtjtjjj||d\} } } Wdn1swYn+tjdtjd<tjjj||d\} } } dtjd<Wdn1swYt| ||S)a Loading a PaddlePaddle Model and Exporting the TensorRT-Optimized Program. Args: model_path(str):The directory path where the PaddlePaddle model is located. The model path can either include the model directory and prefix (e.g., 'model_dir/inference'), or it can be the full path to the model (e.g., 'model_dir/inference.json'). config(TensorRTConfig):The configuration of TensorRTConfig. Returns: program:The TensorRT optimized program. Examples: .. code-block:: python >>> # example 1: >>> # This example takes the user-specified model input shape, and Paddle internally generates corresponding random data. >>> import numpy as np >>> import paddle >>> import paddle.inference as paddle_infer >>> import paddle.nn.functional as F >>> from paddle import nn >>> from paddle.tensorrt.export import Input, TensorRTConfig >>> class LinearNet(nn.Layer): >>> def __init__(self, input_dim): >>> super().__init__() >>> self.linear = nn.Linear(input_dim, input_dim) >>> def forward(self, x): >>> return F.relu(self.linear(x)) >>> input_dim = 3 >>> # 1.Instantiate the network. >>> layer = LinearNet(input_dim) >>> save_path = "/tmp/linear_net" >>> # 2.Convert dynamic graph to static graph and save as a JSON file. >>> paddle.jit.save(layer, save_path, [paddle.static.InputSpec(shape=[-1, input_dim])]) >>> # 3.Create TensorRTConfig >>> input_config = Input( >>> min_input_shape=[1, input_dim], >>> optim_input_shape=[2, input_dim], >>> max_input_shape=[4, input_dim], >>> name='x', >>> ) >>> trt_config = TensorRTConfig(inputs=[input_config]) >>> trt_config.save_model_dir = "/tmp/linear_net_trt" >>> # 4.Perform TensorRT conversion >>> program_with_trt = paddle.tensorrt.convert(save_path, trt_config) >>> # 5.Create a Predictor and run TensorRT inference. >>> config = paddle_infer.Config( >>> trt_config.save_model_dir + '.json', >>> trt_config.save_model_dir + '.pdiparams', >>> ) >>> config.enable_use_gpu(100, 0) >>> predictor = paddle_infer.create_predictor(config) >>> input_data = np.random.randn(2, 3).astype(np.float32) >>> model_input = paddle.to_tensor(input_data) >>> output_converted = predictor.run([model_input]) >>> # example 2: >>> # In this example, the user specifies the actual input. >>> import numpy as np >>> import paddle >>> import paddle.inference as paddle_infer >>> import paddle.nn.functional as F >>> from paddle import nn >>> from paddle.tensorrt.export import Input, TensorRTConfig >>> class LinearNet(nn.Layer): >>> def __init__(self, input_dim): >>> super().__init__() >>> self.linear = nn.Linear(input_dim, input_dim) >>> def forward(self, x): >>> return F.relu(self.linear(x)) >>> input_dim = 3 >>> # 1.Instantiate the network. >>> layer = LinearNet(input_dim) >>> save_path = "/tmp/linear_net" >>> # 2.Convert dynamic graph to static graph and save as a JSON file. >>> paddle.jit.save(layer, save_path, [paddle.static.InputSpec(shape=[-1, input_dim])]) >>> # 3.Create TensorRTConfig >>> input_config = Input( >>> warmup_data=( >>> np.random.rand(1,3).astype(np.float32), >>> np.random.rand(2,3).astype(np.float32), >>> np.random.rand(4,3).astype(np.float32), >>> ), >>> name='x', >>> ) >>> trt_config = TensorRTConfig(inputs=[input_config]) >>> trt_config.save_model_dir = "/tmp/linear_net_trt" >>> # 4.Perform TensorRT conversion >>> program_with_trt = paddle.tensorrt.convert(save_path, trt_config) >>> # 5.Create a Predictor and run TensorRT inference. >>> config = paddle_infer.Config( >>> trt_config.save_model_dir + '.json', >>> trt_config.save_model_dir + '.pdiparams', >>> ) >>> config.enable_use_gpu(100, 0) >>> predictor = paddle_infer.create_predictor(config) >>> input_data = np.random.randn(2, 3).astype(np.float32) >>> model_input = paddle.to_tensor(input_data) >>> output_converted = predictor.run([model_input]) z{The `config.save_model_dir` and `model_path` cannot be the same. Please specify a different directory for saving the model.rTz.jsonz.pdmodelFzUnsupported extension z. Only support json/pdmodelz .pdiparamsz,No valid model file found in the directory 'ze'. Expected either 'json' or 'pdmodel'. Please ensure that the directory contains one of these files.zParameters file 'zJ' not found. Please ensure the weights file exists in the model directory.)rN1ZFLAGS_enable_pir_in_executor0)ospathabspathrFr#rZrxZ global_scoperwryisfilesplitsplitextjoinexistsrqrrioZload_inference_modelZ OldIrGuardenvironr)Z model_pathrrbr~rZis_jsonZ model_dirZ model_fileZ model_prefixextZ params_pathriZfeed_target_namesZ fetch_targetsr'r'r(convertwsh{              r)NNN)+ __future__rloggingrenumrnumpyr0rZZ paddle.baserrZpaddle.base.executorrZpaddle.base.frameworkrZpaddle.base.log_helperrZpaddle.jit.apir r r Z'paddle.jit.dy2static.program_translatorr Z paddle.nnr Zpaddle.tensorrt.converterrZpaddle.tensorrt.utilrrrrrr8INFOr$rr;r@rrrr'r'r'r(s4        "r _ @