tf.experimental.tensorrt.Converter

Офлайн-конвертер для преобразования TF-TRT для сохранённых моделей TF 2.0.

tf.experimental.tensorrt.Converter(
    input_saved_model_dir=None, input_saved_model_tags=None,
    input_saved_model_signature_key=None, conversion_params=None
)

В настоящее время это недоступно на платформе Windows.

Обратите внимание, что в версии 2 не поддерживается is_dynamic_op=False, что означает, что движки TRT будут построены только при выполнении соответствующего оператора TRTEngineOp. Но мы по-прежнему предоставляем способ избежать затрат на построение движков TRT во время вывода (см. дополнительные сведения ниже).

Существует несколько способов выполнения преобразования:

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)

В этом случае движки TRT не будут построены или сохранены в преобразованной сохранённой модели. Но если данные входных данных доступны во время преобразования, мы всё ещё можем построить и сохранить движки TRT, чтобы уменьшить затраты во время вывода (см. вариант 2 ниже).

1. Точность 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.
# With TensorRT 5.1, different engines will be built (and saved later) for
# different input shapes to the TRTEngineOp.
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() в среде, аналогичной производственной (например, с тем же типом графического процессора).

1. Точность 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)

Методы

build

Просмотреть исходный код

build(
    input_fn
)

Запустить вывод с преобразованным графом для построения движков TensorRT.

convert

Просмотреть исходный код

convert(
    calibration_input_fn=None
)

Преобразовать входную SavedModel в формате 2.0.

save

Просмотреть исходный код

save(
    output_saved_model_dir
)

Сохранить преобразованную SavedModel.