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 引擎相同的配置文件。