Python tf.experimental.tensorrt.Converter用法及代码示例

用于 TF 2.0 SavedModels 的 TF-TRT 转换的离线转换器。

用法

tf.experimental.tensorrt.Converter(
    input_saved_model_dir=None, input_saved_model_tags=None,
    input_saved_model_signature_key=None, use_dynamic_shape=None,
    dynamic_shape_profile_strategy=None,
    max_workspace_size_bytes=DEFAULT_TRT_MAX_WORKSPACE_SIZE_BYTES,
    precision_mode=TrtPrecisionMode.FP32, minimum_segment_size=3,
    maximum_cached_engines=1, use_calibration=True, allow_build_at_runtime=True,
    conversion_params=None
)

目前这在 Windows 平台上不可用。

有几种方法可以运行转换：

1. FP32/FP16精度

   params = tf.experimental.tensorrt.ConversionParams(
       precision_mode='FP16')
   converter = tf.experimental.tensorrt.Converter(
       input_saved_model_dir="my_dir", conversion_params=params)
   converter.convert()
   converter.save(output_saved_model_dir)

   在这种情况下，不会在转换后的 SavedModel 中构建或保存任何 TRT 引擎。但是如果在转换过程中输入数据可用，我们仍然可以构建和保存 TRT 引擎以降低推理过程中的成本(参见下面的选项 2)。

2. FP32/FP16 精度与预建引擎

   params = tf.experimental.tensorrt.ConversionParams(
       precision_mode='FP16',
       # Set this to a large enough number so it can cache all the engines.
       maximum_cached_engines=16)
   converter = tf.experimental.tensorrt.Converter(
       input_saved_model_dir="my_dir", conversion_params=params)
   converter.convert()
   
   # Define a generator function that yields input data, and use it to execute
   # the graph to build TRT engines.
   def my_input_fn():
     for _ in range(num_runs):
       inp1, inp2 = ...
       yield inp1, inp2
   
   converter.build(input_fn=my_input_fn)  # Generate corresponding TRT engines
   converter.save(output_saved_model_dir)  # Generated engines will be saved.

   这样，将为 TRTEngineOp 的每个唯一输入形状构建/保存一个引擎。这对于无法在推理期间构建引擎但可以访问与生产中使用的数据相似的输入数据(例如，具有相同输入形状)的应用程序很有用。此外，生成的 TRT 引擎依赖于平台，因此我们需要在与生产类似的环境中运行build()(例如，使用相同类型的 GPU)。

3. INT8 精度和预置引擎校准

   params = tf.experimental.tensorrt.ConversionParams(
       precision_mode='INT8',
       # Currently only one INT8 engine is supported in this mode.
       maximum_cached_engines=1,
       use_calibration=True)
   converter = tf.experimental.tensorrt.Converter(
       input_saved_model_dir="my_dir", conversion_params=params)
   
   # Define a generator function that yields input data, and run INT8
   # calibration with the data. All input data should have the same shape.
   # At the end of convert(), the calibration stats (e.g. range information)
   # will be saved and can be used to generate more TRT engines with different
   # shapes. Also, one TRT engine will be generated (with the same shape as
   # the calibration data) for save later.
   def my_calibration_input_fn():
     for _ in range(num_runs):
       inp1, inp2 = ...
       yield inp1, inp2
   
   converter.convert(calibration_input_fn=my_calibration_input_fn)
   
   # (Optional) Generate more TRT engines offline (same as the previous
   # option), to avoid the cost of generating them during inference.
   def my_input_fn():
     for _ in range(num_runs):
       inp1, inp2 = ...
       yield inp1, inp2
   converter.build(input_fn=my_input_fn)
   
   # Save the TRT engine and the engines.
   converter.save(output_saved_model_dir)

4. 要使用动态形状，我们需要使用输入函数调用 build 方法来生成配置文件。此步骤类似于上述 INT8 校准步骤。转换器还需要使用 use_dynamic_shape=True 和以下 profile_strategies 之一创建，以便根据输入函数产生的输入创建配置文件：

   • Range：创建一个适用于维度值在 [min_dims, max_dims] 范围内的输入的配置文件，其中 min_dims 和 max_dims 是从提供的输入派生的。
   • Optimal ：为每个输入创建一个配置文件。该配置文件仅适用于与其创建的输入具有相同维度的输入。 GPU 引擎将以此类输入的最佳性能运行。
   • Range+Optimal ：为 Range 和 Optimal 创建配置文件。
   • ImplicitBatchModeCompatible ：创建将产生与 implicit_batch_mode 将产生的 GPU 引擎相同的配置文件。