Запуск нескольких оптимизированных моделей TensorRT в Tensorflow - PullRequest
Новая
       22

Запуск нескольких оптимизированных моделей TensorRT в Tensorflow

1 голос
/ 01 апреля 2020

Мой проект использует несколько моделей Keras. Эти модели могут иметь вход с разным размером пакета, который варьируется от 1 до 24. Я решил оптимизировать эти модели, используя TF-TRT.

Я пробовал 2 подхода к конвертации: 

from tensorflow.python.compiler.tensorrt import trt_convert as trt

Первый подход преобразует модель, но не создает механизмы TensorRT для модели: 

conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
                precision_mode=trt.TrtPrecisionMode.FP32)
converter = trt.TrtGraphConverterV2(
                input_saved_model_dir=saved_model_path,
                conversion_params=conversion_params)
converter.convert()
converter.save(output_saved_model_dir=trt_fp32_model_path)

Второй подход преобразует модель и создает механизм TensorRT для всех возможных форм ввода: 

def input_function():
    def input_function():
    input_shapes = [(x, MODEL_INPUT_H, MODEL_INPUT_W, 3) for x in range(1, 25)]
    for shape in input_shapes:
         yield [np.random.normal(size=shape).astype(np.float32)]

conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(
    precision_mode=trt.TrtPrecisionMode.FP32,
    maximum_cached_engines=100
)

converter = trt.TrtGraphConverterV2(
    input_saved_model_dir=saved_model_path,
    conversion_params=conversion_params)

converter.convert()
converter.build(input_fn=input_function)
converter.save(output_saved_model_dir=trt_fp32_model_path)

В сценарии, использующем мои модели, я последовательно использую эти модели: 

some loop:
    model1.predict(model1_input)
    model2.predict(model2_input)
    model3.predict(model3_input)

Когда для оптимизации моделей используется первый подход к преобразованию, я могу загрузить все модели, но во время выполнения Tensorflow перестраивает движки TensorRT каждые время изменения контекста модели. Это приводит к значительному снижению производительности, которое я пытался преодолеть, кэшируя механизмы TensorRT для этих моделей (второй подход к преобразованию).

Проблема заключается в том, что при попытке загрузить более одной оптимизированной модели TensorRT с предустановленным -встроенных движках, Tensorflow выдает следующую ошибку: 

2020-04-01 09:11:44.820866: W tensorflow/core/common_runtime/base_collective_executor.cc:216] BaseCollectiveExecutor::StartAbort Internal: Expect engine cache to be empty, but got 24 entries.
     [[{{node StatefulPartitionedCall/InitializeTRTResource}}]]
Error -  Expect engine cache to be empty, but got 24 entries.
     [[{{node StatefulPartitionedCall/InitializeTRTResource}}]] [Op:__inference_restored_function_body_64832]

Function call stack:
restored_function_body

Та же самая ошибка возникает, когда для каждой модели сохраняется только один движок.

Я использую следующий код для загрузки оптимизированной для TensorRT SavedModel: 

saved_model_loaded = tf.saved_model.load(
            trt_fp32_model_path,
            tags=[tag_constants.SERVING]
        )
graph_func = saved_model_loaded.signatures['serving_default']

Я также пытался преобразовать graph_fun c в frozen_fun c, но это не имело никакого значения: 

graph_func = saved_model_loaded.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

frozen_func = convert_to_constants.convert_variables_to_constants_v2(
     graph_func)

Я использую nvcr.io/nvidia/tensorflow:19.12-tf2-py3 docker контейнер для оптимизации / запуска моделей.

Можно ли вообще запустить одновременно несколько оптимизированных TensorRT моделей с предварительно собранными двигателями, использующими Tensorflow? Или это можно сделать только с помощью сервера вывода TensorRT?

Если это допустимый сценарий использования, что мне не хватает в моем рабочем процессе?

...