使用 TensorRT 加速 TensorFlow 中的推理用户指南

摘要

TensorFlow-TensorRT (TF-TRT) 是一个用于 TensorFlow 的深度学习编译器，可优化 TF 模型以在 NVIDIA 设备上进行推理。TF-TRT 是 NVIDIA TensorRT (TRT) 高性能深度学习推理 SDK 的 TensorFlow 集成，允许用户直接在 TensorFlow 框架内利用其功能。

TF-TRT 根据与 TensorRT 的兼容性自动将 TensorFlow 图划分为子图。这些兼容的子图由 TensorRT 优化和执行，其余图的执行则委托给原生 TensorFlow。这允许使用 TensorFlow 丰富的功能集，同时尽可能使用 TensorRT 优化图，从而提供灵活性和性能。要了解有关 TensorRT 提供的优化的更多信息，请访问 TensorRT 页面。

使用 TF-TRT 最重要的好处是，用户可以在 TensorFlow 上创建和测试他们的模型，并利用 TensorRT 提供的性能加速，只需添加几行代码，而无需直接使用 TensorRT 在 C++ 中进行开发。这实现了从模型定义、训练到在 NVIDIA 设备上部署的无缝工作流程。

对于需要完全控制 TensorRT 功能的专家用户，建议将 TensorFlow 模型导出到 ONNX 并直接使用 TensorRT。

TF-TRT 通过其 Python 或 C++ API 摄取从训练后的 TensorFlow 模型创建的 TensorFlow SavedModel（参见 构建和加载 SavedModel）。支持的子图被替换为 TensorRT 优化的节点（称为 TRTEngineOp），从而生成一个新的 TensorFlow 图，该图将同时具有 TensorFlow 和 TensorRT 组件，如下图所示

在将子图转换为 TRTEngineOp 的过程中，TensorRT 对神经网络图执行多个重要的转换和优化，包括常量折叠、剪枝不必要的图节点、层融合等等。有关优化的完整列表，请参阅 TensorRT 文档。

转换为单个 TensorRT 引擎的操作越多，使用 TensorRT 获得的潜在好处就越大。因此，TF-TRT 的目标是转换 TensorFlow 计算图中的最大操作数，同时限制 TensorRT 引擎的总数，从而最大限度地提高性能。根据图中的操作，最终图可能具有多个 TensorRT 节点。请参阅 TensorRT 的文档，以了解有关 特定图优化的更多信息。

如果将段转换为 TensorRT 引擎失败或执行生成的 TensorRT 引擎失败，则 TFTRT 将尝试执行原生 TensorFlow 段。这称为原生段回退。

NVIDIA TensorFlow 的 NGC 容器 在构建和测试时启用了 TF-TRT 支持，允许在容器中开箱即用，而无需费力设置自定义环境。要了解如何拉取容器，请查看 NGC 用户指南，但是可以使用简单的 docker 命令下载它们

            

            
docker pull nvcr.io/nvidia/tensorflow:<yy.mm>-tf2-py3
# eg: The TF2 container from January 2022 (22.01) would be
docker pull nvcr.io/nvidia/tensorflow:22.01-tf2-py3
        

TF-TRT 也可在 TensorFlow 存储库中找到，并且可以与任何附带 TensorRT 的 TensorFlow 安装一起使用。这意味着当您安装 tensorflow-gpu 时，它包含 TF-TRT，TensorRT 安装后即可直接使用。

要从头开始构建，请按照这些 说明 进行操作。您需要在 bazel 配置中启用 TensorRT（默认情况下禁用）。

容器中的 nvidia-examples 目录下还提供了示例，可以执行这些示例来测试 TF-TRT。某些模型可能需要其他软件包才能正确执行。

为了使用 TF-TRT 优化 Tensorflow 模型，tensorflow 模型需要以 SavedModel 格式保存。使用 TF-TRT 进行转换只是优化 NVIDIA 设备上推理模型的一个额外步骤。

TF-TRT 工作流程很简单。基本步骤如下

1. 使用 TF-TRT 转换器转换训练后的 SavedModel。
2. 构建 TRT 引擎（可选，运行模型将自动构建它）。
3. 保存转换后的模型以供将来使用（可选）。
4. 完成！像往常一样使用转换后的模型运行推理。

以下是一个完整的 Python 示例，从模型定义和训练到 TF-TRT 转换和推理。

1. 使用 Tensorflow 创建您的模型。为了说明目的，我们将使用 Keras 顺序模型（来自 此处）。
   复制

   已复制！

               
   
               
   import tensorflow as tf
   from tensorflow import keras
    
   # Define a simple sequential model
   model = tf.keras.models.Sequential([
      tf.keras.layers.Flatten(input_shape=(28, 28)),
      tf.keras.layers.Dense(128, activation='relu'),
      tf.keras.layers.Dropout(0.2),
      tf.keras.layers.Dense(10)
   ])
    
   model.compile(optimizer='adam',
      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
      metrics=['accuracy'])
           

2. 获取数据集并根据需要对其进行预处理。
   复制

   已复制！

               
   
               
   mnist = tf.keras.datasets.mnist
    
   (x_train, y_train), (x_test, y_test) = mnist.load_data()
   x_train, x_test = x_train / 255.0, x_test / 255.0
   x_train = tf.cast(x_train, dtype=tf.float32)
   y_train = tf.cast(y_train, dtype=tf.float32)
   x_test = tf.cast(x_test, dtype=tf.float32)
   y_test = tf.cast(y_test, dtype=tf.float32)
           

3. 根据需要训练模型并评估准确率，如 此处 所示。
   复制

   已复制！

               
   
               
   # Train the model
   model.fit(x_train, y_train, epochs=5)
    
   # Evaluate your model accuracy
   model.evaluate(x_test,  y_test, verbose=2)
           

4. 以 saved_model 格式保存模型。
   复制

   已复制！

               
   
               
   # Save model in the saved_model format
   SAVED_MODEL_DIR="./models/native_saved_model"
   model.save(SAVED_MODEL_DIR)
           

5. 使用 TF-TRT 转换器转换模型。
   复制

   已复制！

               
   
               
   from tensorflow.python.compiler.tensorrt import trt_convert as trt
    
   # Instantiate the TF-TRT converter
   converter = trt.TrtGraphConverterV2(
      input_saved_model_dir=SAVED_MODEL_DIR,
      precision_mode=trt.TrtPrecisionMode.FP32
   )
    
   # Convert the model into TRT compatible segments
   trt_func = converter.convert()
   converter.summary()
           

6. 构建 TRT 引擎。
   复制

   已复制！

               
   
               
   MAX_BATCH_SIZE=128
   def input_fn():
      batch_size = MAX_BATCH_SIZE
      x = x_test[0:batch_size, :]
      yield [x]
    
   converter.build(input_fn=input_fn)
           

7. 保存转换后的模型以供将来使用。
   复制

   已复制！

               
   
               
   OUTPUT_SAVED_MODEL_DIR="./models/tftrt_saved_model"
   converter.save(output_saved_model_dir=OUTPUT_SAVED_MODEL_DIR)
           

8. 使用转换后的模型运行推理。
   复制

   已复制！

               
   
               
   # Get batches of test data and run inference through them
   infer_batch_size = MAX_BATCH_SIZE // 2
   for i in range(10):
      print(f"Step: {i}")
     
      start_idx = i * infer_batch_size
      end_idx   = (i + 1) * infer_batch_size
      x = x_test[start_idx:end_idx, :]
    
      trt_func(x)
           

有关更多示例，请参阅 TensorRT 存储库中的 TF-TRT 示例。

TF-TRT 利用 TensorRT 的许多功能来加速推理。其中一些功能包括

• 混合精度执行（FP32、FP16 和 INT8）
• INT8 量化
• 动态批处理和输入形状

本节将概述上述功能，并提供有关如何利用它们的用法/最佳实践示例。

4.1. 支持的精度级别

TensorRT 可以将张量和权重转换为较低的精度，以便在优化期间实现更快的推理。precision_mode 参数设置精度模式；可以是 FP32、FP16 或 INT8 之一。低于 FP32 的精度（例如 FP16 和 INT8）可以从 TensorRT 引擎中提取更高的性能。如果可能，FP16 模式使用 Tensor Core 或半精度硬件指令。INT8 精度模式使用 整数硬件指令。

鼓励用户尝试降低的精度模式，例如 FP16 和 INT8。FP16 将提高性能，而不会造成明显的精度损失；使用 AMP 训练的模型应该没有损失。INT8 精度模式将具有最佳性能，但是，INT8 量化引起的量化误差可能会在某些模型中引入精度下降。请参阅 量化 部分，以更详细地了解这种影响以及如何缓解它。

应根据应用程序要求（性能、内存消耗、准确率）选择精度级别。无论选择如何，始终建议在转换为 TensorRT 后进行模型验证。

用户应注意，选择较低的精度模式并不意味着整个网络都将以该精度运行。TensorRT 选择所选精度或更高精度中最快的层以获得最佳性能（请参阅 降低精度）。

4.2. 量化

深度学习中的量化是指转换深度学习模型的参数以较低的精度执行计算。这是一种流行的优化方法，有助于减小深度学习模型的大小，从而加快推理速度并降低功耗。这在所有部署中都很有用，但对于在计算能力较低的嵌入式设备（例如 NVIDIA Jetson）上部署至关重要。

为了用示例说明量化，假设将 3.999x2.999 和 4x3 相乘。后者（整数量化）运算比前者快得多。这就是通过将数字量化为较低精度而努力实现的速度提升。简而言之，量化是将输入值从大范围和精细粒度映射到较小范围和粗粒度的输出值的过程，从而降低精度。

在深度学习的背景下，我们通常使用浮点 32 位表示 (FP32) 来训练深度学习模型，因为我们可以利用更广泛的数字范围。在模型量化期间，模型数据（网络参数和激活）从这种浮点表示转换为较低精度的表示，通常使用 8 位整数 (INT8)。不幸的是，这种近似可能会导致模型精度降低。

TF-TRT 中使用的主要量化方法是训练后量化 (PTQ)。顾名思义，训练后量化是一种在先前训练的模型上使用的技术，用于减小模型的大小并获得吞吐量优势，同时减轻模型精度的成本。

由于小的舍入误差会通过网络传播，并且对模型精度产生越来越大的影响，因此已经开发了不同的量化技术，例如量化感知训练 (QAT)，以减轻这种影响。QAT 目前在 TF-TRT 中是实验性的。

4.2.1. 训练后量化

训练后量化（又名 PTQ）在 TensorRT 的上下文中称为 INT8-calibration。

在校准阶段，TensorRT 使用提供的输入“校准”数据来估计网络每个张量的最佳比例和偏差值，给定其动态范围和值分布。TF-TRT 将收集到的此信息存储在转换后的模型中。使用 PTQ 的 TF-TRT 工作流程非常简单。用户需要执行以下操作

• 在转换器实例化期间
  • 将精度模式设置为 INT8。
  • 将 use_calibration 标志设置为 True。
• 在转换过程中，用户需要传入代表性的输入数据加载器以进行校准。重要的是，校准输入数据代表模型预期在其上运行的输入范围，以便校准为激活生成有意义的比例因子；数据越多，量化就越准确。例如，测试数据集（或其某些子集）通常是一个好的数据源。

用户应注意，校准数据始终应具有单一形状。

以下 Python 示例演示了校准

            

            
from tensorflow.python.compiler.tensorrt import trt_convert as trt
 
# Instantiate the TF-TRT converter
converter = trt.TrtGraphConverterV2(
   input_saved_model_dir=SAVED_MODEL_DIR,
   precision_mode=trt.TrtPrecisionMode.INT8,
   use_calibration=True
)
 
# Use data from the test/validation set to perform INT8 calibration
BATCH_SIZE=32
NUM_CALIB_BATCHES=10
def calibration_input_fn():
   for i in range(NUM_CALIB_BATCHES):
       start_idx = i * BATCH_SIZE
       end_idx = (i + 1) * BATCH_SIZE
       x = x_test[start_idx:end_idx, :]
       yield [x]
 
# Convert the model with valid calibration data
func = converter.convert(calibration_input_fn=calibration_input_fn)
 
# Input for dynamic shapes profile generation
MAX_BATCH_SIZE=128
def input_fn():
   batch_size = MAX_BATCH_SIZE
   x = x_test[0:batch_size, :]
   yield [x]
 
# Build the engine
converter.build(input_fn=input_fn)
 
OUTPUT_SAVED_MODEL_DIR="./models/tftrt_saved_model"
converter.save(output_saved_model_dir=OUTPUT_SAVED_MODEL_DIR)
 
converter.summary()
 
# Run some inferences!
for step in range(10):
   start_idx = step * BATCH_SIZE
   end_idx   = (step + 1) * BATCH_SIZE
 
   print(f"Step: {step}")
   x = x_test[start_idx:end_idx, :]
   func(x)
        

4.3. 动态形状

TensorFlow 模型可以具有固定或动态形状的输入张量。在训练之前，我们通常将大多数输入维度设置为固定值。例如，一个模型采用一批 8 个输入图像，分辨率为 224x224，颜色通道为 3，可以具有形状为 [8, 224, 224, 3] 的输入张量。根据模型架构，我们可以将某些输入维度保持为动态，以允许使用更广泛的输入形状进行推理：动态输入形状的典型示例包括

• 批处理大小 – 例如，对于图像分类模型，网络输入张量可以是 [?, 224, 224, 3]，其中批处理大小在模型定义期间未知，并且允许在运行时采用不同的值。
• 完全卷积网络的图像大小 [8, ?, ?, 3]
• Transformer 模型的序列长度。例如，BERT 编码器具有形状为 [N, S] 的输入张量，其中 N 是批处理大小，S 是序列长度，并且这两个维度都可以是动态的。

TF-TRT 通过用户提供的有关转换后的模型应支持的输入形状范围的信息来支持具有动态形状的模型。

默认情况下，TF-TRT 允许动态批处理大小。最大批处理大小 (N) 设置为用于为转换后的模型构建引擎的批处理大小。这样的模型将支持 [1..N] 之间的任何批处理大小。（也称为隐式批处理模式）。如果我们尝试使用更大的批处理大小来推断模型，则 TF-TRT 将构建另一个引擎来执行此操作。这会对性能产生重大影响，因为引擎构建成本很高。allow_build_at_runtime 和 max_cached engines 转换参数控制 TF-TRT 的运行时引擎构建行为。

为了支持批处理维度以外的动态输入维度，我们需要通过将 use_dynamic_shape=True 参数传递给转换器来启用动态形状模式。TF-TRT 中的动态形状模式利用 TensorRT 的 动态形状 功能来提高网络的转换率并有效地处理具有未知输入形状的网络。转换率的提高意味着可以在 TensorRT 中运行更多的网络。这提高了此类网络与 TF-TRT 一起使用时的性能。

除了启用 use_dynamic_shape 标志外，还需要向 TF-TRT 提供有关推理期间预期的形状范围的信息，如以下 示例 中所示。

            

            
# Instantiate the TF-TRT converter
# Instantiate the TF-TRT converter
PROFILE_STRATEGY="Optimal"
converter = trt.TrtGraphConverterV2(
   input_saved_model_dir=bert_saved_model_path,
   precision_mode=trt.TrtPrecisionMode.FP32,
   use_dynamic_shape=True,
   dynamic_shape_profile_strategy=PROFILE_STRATEGY)
 
# Convert the model to TF-TRT
converter.convert()
 
VOCAB_SIZE = 30522  # Model specific, look in the model README.
# Build engines for input sequence lengths of 128, and 384.
input_shapes = [[(1, 128), (1, 128), (1, 128)],
               [(1, 384), (1, 384), (1, 384)]]
def input_fn():
   for shapes in input_shapes:
       # return a list of input tensors
       yield [tf.convert_to_tensor(
         np.random.randint(low=0, high=VOCAB_SIZE, size=x,dtype=np.int32))
         for x in shapes]
 
converter.build(input_fn)
        

在保存转换后的模型之前，需要构建它以通过使用 input_fn 来处理一定范围的输入参数。与校准输入不同，这些输入不需要表示真实输入数据，对于大多数模型，只有输入形状很重要；数据相关的形状是例外。

上面的示例说明了一个类似 BERT 的模型，它有三个输入张量。我们的 input_fn 定义了两种不同的输入形状，一种序列长度为 128，另一种序列长度为 384。

可以使用 dynamic_shape_profile_strategy 参数进一步指定动态输入。此参数选择用于定义 TensorRT 的优化配置文件的策略（其中“优化配置文件”是 TensorRT 用于描述输入形状信息的术语）。以下是优化配置文件的选项

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

下图和表格说明了配置文件策略如何影响转换后的模型接受的形状范围。

如果仅预期少量具体输入形状，则建议使用“最优”策略。

如果未调用 build()，则 TensoRT 引擎创建将在首次推断转换后的模型时进行。在此推理期间使用的输入形状会将 TensorRT 配置文件策略设置为默认策略“范围”，参数为 min_dims=max_dims。

4.4. 简单示例

4.4.1. Python 示例

以下是一个简单的 Python 示例，演示了使用随机输入转换 BERT 模型。

            

            
# Prerequisite: Install the python module below before running this example.
# pip install -q tf-models-official
 
import tensorflow as tf
import tensorflow_hub as hub
 
tfhub_handle_encoder = 'https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/3'
bert_saved_model_path = './models/bert_base'
 
bert_model = hub.load(tfhub_handle_encoder)
tf.saved_model.save(bert_model, bert_saved_model_path)
 
import numpy as np
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.compiler.tensorrt import trt_convert as trt
 
# Instantiate the TF-TRT converter
PROFILE_STRATEGY="Optimal"
converter = trt.TrtGraphConverterV2(
   input_saved_model_dir=bert_saved_model_path,
   precision_mode=trt.TrtPrecisionMode.FP32,
   use_dynamic_shape=True,
   dynamic_shape_profile_strategy=PROFILE_STRATEGY)
 
# Convert the model to TF-TRT
converter.convert()
 
VOCAB_SIZE = 30522  # Model specific, look in the model README.
# Build engines for input sequence lengths of 128, and 384.
input_shapes = [[(1, 128), (1, 128), (1, 128)],
               [(1, 384), (1, 384), (1, 384)]]
def input_fn():
   for shapes in input_shapes:
       # return a list of input tensors
       yield [tf.convert_to_tensor(
         np.random.randint(low=0, high=VOCAB_SIZE, size=x,dtype=np.int32))
         for x in shapes]
 
converter.build(input_fn)
 
# Save the converted model
bert_trt_path = "./models/tftrt_bert_base"
converter.save(bert_trt_path)
converter.summary()
 
# Some helper functions
def get_func_from_saved_model(saved_model_dir):
   saved_model_loaded = tf.saved_model.load(
       saved_model_dir, tags=[tag_constants.SERVING])
   graph_func = saved_model_loaded.signatures[
       signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
   return graph_func, saved_model_loaded
 
 
def get_random_input(batch_size, seq_length):
   # Generate random input data
   mask = tf.convert_to_tensor(np.ones((batch_size, seq_length), dtype=np.int32))
   type_id = tf.convert_to_tensor(np.zeros((batch_size, seq_length), dtype=np.int32))
   word_id = tf.convert_to_tensor(
     np.random.randint(0, VOCAB_SIZE, size=[batch_size, seq_length], dtype=np.int32))
   return {'input_mask':mask, 'input_type_ids': type_id, 'input_word_ids':word_id}
 
 
# Get a random input tensor
input_tensor = get_random_input(1, 128)
 
# Specify the output tensor interested in. This output is the 'classifier'
result_key = 'bert_encoder_1'
trt_func, _ = get_func_from_saved_model(bert_trt_path)
 
## Let's run some inferences!
for i in range(0, 10):
   print(f"Step: {i}")
   preds = trt_func(**input_tensor)
   result = preds[result_key]
        

4.4.2. C++ 示例

Tensorflow 支持在 C++ 中执行模型，包括 TF-TRT 转换后的模型。TF-TRT 转换器的 C++ API 目前是实验性的。

首先通过 TF-TRT Python API 转换模型。

下面针对具有合成数据的模型解释了加载和运行模型的 C++ 工作流程。

1. 初始化 Tensorflow 所需的全局状态
2. 加载保存的模型并初始化 TF 会话
3. 设置输入
4. 设置输出
5. 运行推理
6. 释放资源

显示上述步骤的示例

1. 初始化 Tensorflow 和 TF 会话所需的全局状态。
   复制

   已复制！

               
   
               
   // We need to call this to set up global state for TensorFlow.
    tensorflow::port::InitMain(argv[0], &argc, &argv);
    if (argc > 1) {
      LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage;
      return -1;
    }
           

2. 加载保存的模型并初始化 TF 会话。
   复制

   已复制！

               
   
               
   // Some helper functions
   // Returns info for nodes listed in the signature definition.
   std::vector<tensorflow::TensorInfo> GetNodeInfo(
      const google::protobuf::Map<string, tensorflow::TensorInfo>& signature) {
    std::vector<tensorflow::TensorInfo> info;
    for (const auto& item : signature) {
      info.push_back(item.second);
    }
    return info;
   }
    
   // Load the `SavedModel` located at `model_dir`.
   Status LoadModel(const string& model_dir, const string& signature_key,
                   tensorflow::SavedModelBundle* bundle,
                   std::vector<tensorflow::TensorInfo>* input_info,
                   std::vector<tensorflow::TensorInfo>* output_info) {
    tensorflow::RunOptions run_options;
    tensorflow::SessionOptions sess_options;
    
    tensorflow::OptimizerOptions* optimizer_options =
        sess_options.config.mutable_graph_options()->mutable_optimizer_options();
    optimizer_options->set_opt_level(tensorflow::OptimizerOptions::L0);
    optimizer_options->set_global_jit_level(tensorflow::OptimizerOptions::OFF);
    
    sess_options.config.mutable_gpu_options()->force_gpu_compatible();
    TF_RETURN_IF_ERROR(tensorflow::LoadSavedModel(sess_options, run_options,
                                                  model_dir, {"serve"}, bundle));
    
    // Get input and output names
    auto signature_map = bundle->GetSignatures();
    const tensorflow::SignatureDef& signature = signature_map[signature_key];
    *input_info = GetNodeInfo(signature.inputs());
    *output_info = GetNodeInfo(signature.outputs());
    
    return Status::OK();
   }
           

3. 设置输入。在这里，我们使用合成数据并将其放置在设备上以进行推理。
   复制

   已复制！

               
   
               
   // Create random inputs matching `input_info`
   Status SetupInputs(int32_t batch_size,
                     int32_t input_size,
                     std::vector<tensorflow::TensorInfo>& input_info,
                     std::vector<std::pair<std::string, tensorflow::Tensor>>* inputs) {
    
    //std::vector<std::pair<std::string, tensorflow::Tensor>> input_tensors;
    for (auto& info : input_info) {
        // Set input batch size
        auto* shape = info.mutable_tensor_shape();
        shape->mutable_dim(0)->set_size(batch_size);
    
        // Set dynamic dims to static size
        for (size_t i = 1; i < shape->dim_size(); i++) {
            auto* dim = shape->mutable_dim(i);
            if (dim->size() < 0) {
              dim->set_size(input_size);
            }
        }
       
        // Allocate memory and fill host tensor
        Tensor input_tensor(info.dtype(), *shape);
        std::fill_n((uint8_t*)input_tensor.data(), input_tensor.AllocatedBytes(), 1);
    
        inputs->push_back({info.name(), input_tensor});
    }
    
      return Status::OK();
   }
           

4. 设置输出。
   复制

   已复制！

               
   
               
   // Get output tensor names based on `output_info`.
   Status SetupOutputs(std::vector<tensorflow::TensorInfo>& output_info,
                     std::vector<string>* output_names,
                     std::vector<Tensor>* outputs) {
    for (auto& info : output_info) {
      output_names->push_back(info.name());
      outputs->push_back({});
    }
    return Status::OK();
   }
           

5. 运行推理。
   复制

   已复制！

               
   
               
   // Setup inputs
    std::vector<std::pair<std::string, tensorflow::Tensor>> inputs;
    TFTRT_ENSURE_OK(SetupInputs(batch_size, input_size, input_info, &inputs));
    
    // Setup outputs
    std::vector<string> output_names;
    std::vector<Tensor> outputs;
    TFTRT_ENSURE_OK(SetupOutputs(output_info, &output_names, &outputs));
    
    int num_iterations = 10;
    for (int i = 0; i < num_iterations; i++) {
      LOG(INFO) << "Step: " << i;
    
      TFTRT_ENSURE_OK(
          bundle.session->Run(inputs, output_names, {}, &outputs));
    }
           

这是一个将上述所有功能联系在一起的程序

            

            
int main(int argc, char* argv[]) {
 // Parse arguments
 string model_path = "/path/to/model/";
 string signature_key = "serving_default";
 int32_t batch_size = 64;
 int32_t input_size = 128;
 std::vector<Flag> flag_list = {
     Flag("model_path", &model_path, "graph to be executed"),
     Flag("signature_key", &signature_key, "the serving signature to use"),
     Flag("batch_size", &batch_size, "batch size to use for inference"),
     Flag("input_size", &input_size, "shape to use for -1 input dims"),
 };
 string usage = tensorflow::Flags::Usage(argv[0], flag_list);
 const bool parse_result = tensorflow::Flags::Parse(&argc, argv, flag_list);
 if (!parse_result) {
   LOG(ERROR) << usage;
   return -1;
 }
 
 // We need to call this to set up global state for TensorFlow.
 tensorflow::port::InitMain(argv[0], &argc, &argv);
 if (argc > 1) {
   LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage;
   return -1;
 }
 
 // Setup TF session
 tensorflow::SavedModelBundle bundle;
 std::vector<tensorflow::TensorInfo> input_info;
 std::vector<tensorflow::TensorInfo> output_info;
 TFTRT_ENSURE_OK(
     LoadModel(model_path, signature_key, &bundle, &input_info, &output_info));
 
 // Setup inputs
 std::vector<std::pair<std::string, tensorflow::Tensor>> inputs;
 TFTRT_ENSURE_OK(SetupInputs(batch_size, input_size, input_info, &inputs));
 
 // Setup outputs
 std::vector<string> output_names;
 std::vector<Tensor> outputs;
 TFTRT_ENSURE_OK(SetupOutputs(output_info, &output_names, &outputs));
 
 int num_iterations = 10;
 for (int i = 0; i < num_iterations; i++) {
   LOG(INFO) << "Step: " << i;
 
   TFTRT_ENSURE_OK(
       bundle.session->Run(inputs, output_names, {}, &outputs));
 }
 
 return 0;
}
        

此处提供了更多处理真实数据的 C++ 示例。

由 TensorFlow-TensorRT 加速的模型可以使用 NVIDIA Triton 推理服务器进行服务，这是一个开源推理服务软件，可帮助标准化模型部署和执行，并在生产中提供快速且可扩展的 AI。

浏览此 示例 和 NVIDIA Triton 推理服务器 Github 存储库 以获取更多信息！

6.1. 最小段大小

由 TensorRT 优化的 TensorFlow 子图包含一定数量的运算符。如果子图中包含的运算符数量太少，则与执行原始子图相比，为此子图启动 TensorRT 引擎可能效率不高。您可以使用参数 minimum_segment_size 控制子图的大小。将此值设置为 x（默认为 3）将不会为由少于 x 个节点组成的子图生成 TensorRT 引擎。

使用 minimum_segment_size 参数允许您减轻和避免这些小型 TensorRT 引擎引入的潜在开销，并且还可以解决这些引擎引起的任何可能的错误。

minimum_segment_size 的最佳值是特定于模型的。要找到最佳值，请查看转换器摘要（请参阅部分），观察引擎的数量及其大小，并设置一个最大化转换节点数量的值，同时限制子图（TRTEngineOps）的数量。

一个好的规则是从 minimum_segment_size=3 开始，仅在 TF-TRT 生成过多引擎时才更新：此数字会有所不同，但最多 5-10 个引擎是一个好的启发式方法。逐步更新该值，直到引擎数量降至合理的水平，而不会产生性能损失。用于调整引擎数量的一个好的启发式方法是引擎中的节点数。如果段中的节点数与其他段相比太小，则可以将最小段大小增加到大于该段大小的值，以避免转换该段。

6.2. 最大工作区大小

除了用于权重和激活的内存之外，某些 TensorRT 算法还需要临时工作区。max_workspace_size_bytes 参数限制了 TensorRT 可以用于工作区的最大大小。工作区也通过 TensorFlow 分配器分配。如果该值太小，TensorRT 将无法使用某些需要大工作区的算法，从而导致引擎构建失败和原生段回退。

尽管 TensorRT 被允许使用最多需要 max_workspace_size_bytes 工作区量的算法，但所有 TensorRT 算法的最大工作区需求可能仍然小于 max_workspace_size_bytes（这意味着，TensorRT 可能没有任何需要那么多工作区的算法）。在这种情况下，TensorRT 只会分配所需的工作区，而不是分配用户指定的数量。

TensorRT 在 TensorRT 8.4 中引入了延迟分配，如果需要，可以使用给定设备的总全局内存大小。因此，建议使用模型所需的最大值，仅当要在单个设备上构建多个引擎时才限制此大小。

6.3. 转换报告

您可以使用 converter.summary() 查看转换后的引擎列表以及相关信息，例如具体的输入形状和类型。如果您关心未转换的引擎，请改用转换报告，其中显示了未转换的 ops 类型以及此类 ops 存在的实例数。本节详细讨论了这些诊断工具中的每一个。

6.3.1. converter.summary()

了解 TF-TRT 转换过程的巧妙 API 是摘要 API。可以按如下所示调用它

            

            
converter.convert()
converter.summary()
        

在转换后调用时，此 API 会打印转换摘要。它包括引擎名称、每个引擎的节点数、输入和输出数据类型以及每个 TRTEngineOp 的输入形状等信息。详细报告还包括引擎中 ops 的数量和类型。下面显示了一个示例摘要报告，该报告不仅显示了引擎的数量及其组成 ops，还显示了转换的 ops 总数，以了解完整性。

            

            
TRTEngineOP Name                  Device           # Nodes  # Inputs      # Outputs     Input DTypes           Output Dtypes          Input Shapes           Output Shapes        
========================================================================================================================================================================================
 
----------------------------------
 
TRTEngineOp_000_000               device:GPU:0     9        8             8             ['float32', 'float ... ['float32', 'float ... [[-1, -1, 3], [-1, ... [[1, -1, -1, 3], [ ...
 
   - Const: 1x
   - ExpandDims: 8x
 
----------------------------------
 
TRTEngineOp_000_001               device:GPU:0     919      8             2             ['float32', 'float ... ['float32', 'float32'] [[1, 640, 640, 3], ... [[8, 51150, 90], [ ...
 
   - AddV2: 25x
   - BiasAdd: 13x
   - Cast: 11x
   - ConcatV2: 5x
   - Const: 489x
   - Conv2D: 110x
   - Exp: 2x
   - ExpandDims: 2x
   - FusedBatchNormV3: 97x
   - MaxPool: 4x
   - Mul: 17x
   - Pack: 4x
   - Pad: 1x
   - Relu: 16x
   - Relu6: 77x
   - Reshape: 19x
   - Sigmoid: 1x
   - Slice: 1x
   - Split: 1x
   - Squeeze: 8x
   - Sub: 6x
   - Tile: 5x
   - Transpose: 3x
   - Unpack: 2x
 
======================================================================================================================================================================================
[*] Total number of TensorRT engines: 2
[*] % of OPs Converted: 97.38% [928/953]
        

6.3.2. 转换报告

TF-TRT 现在可以报告由环境变量 TF_TRT_SHOW_DETAILED_REPORT 控制的转换过程的结果。转换报告详细说明了未转换的 ops 的类型和数量，以及它们未能转换的原因。例如，当调试导致大量引擎的转换时，这很有用。

此环境变量可以采用以下值

• TF_TRT_SHOW_DETAILED_REPORT=1 – 这会在 stdout 上打印详细的未转换报告。
• TF_TRT_SHOW_DETAILED_REPORT=非空字符串 – 这会将未转换报告以 CSV 格式导出到环境变量定义的路径。
• 如果未使用此环境变量，则默认报告将打印到 stdout。

以下显示了一个示例未转换报告

            

            
################################################################################
TensorRT unsupported/non-converted OP Report:
   - ResizeBilinear -> 8x
       - [Count: 8x] Cannot Convert Bilinear Resize when align_corners=False
 
   - Identity -> 4x
       - [Count: 4x] excluded by segmenter option. Most likely an input or output node.
 
   - Cast -> 2x
       - [Count: 1x] Cast op is not supported - Allowed input dtypes: [float, half]. Received: int32
       - [Count: 1x] Failed to convert at least one input to a TRT_TensorOrWeights: Unsupported tensorflow data type uint8
 
   - NoOp -> 2x
       - [Count: 2x] Op type NoOp is not supported.
 
   - CombinedNonMaxSuppression -> 1x
       - [Count: 1x] Op type CombinedNonMaxSuppression is not supported.
 
   - Placeholder -> 1x
       - [Count: 1x] excluded by segmenter option. Most likely an input or output node.
 
   - Unpack -> 1x
       - [Count: 1x] The argument `strided_slice_spec` is `absl::nullopt` with `dynamic_input_size_indices` non empty.
 
--------------------------------------------------------------------------------
   - Total nonconverted OPs: 19
   - Total nonconverted OP Types: 7
For more information see http://docs.nvda.net.cn/deeplearning/frameworks/tf-trt-user-guide/index.html#supported-ops.
################################################################################
        

6.4. 日志记录

TF-TRT 在可能的情况下利用 TensorFlow 的日志记录基础设施。可以通过使用环境变量来启用详细日志记录。
调试 TF-TRT 中的失败/回归的一种流行方法是在 TF 中启用详细日志记录，然后运行工作负载。这由环境变量 TF_CPP_VMODULE 控制。此处显示了此环境变量的示例用法

            

            
TF_CPP_VMODULE=trt_logger=2,trt_engine_utils=2,trt_engine_op=2,convert_nodes=2,convert_graph=2,segment=2,trt_shape_optimization_profiles=2,trt_engine_resource_ops=2
        

任何逗号分隔的值都可以省略，变量的值控制日志记录的详细程度。

6.5. 导出 TRT 引擎以进行调试

引擎导出由 TF_TRT_EXPORT_TRT_ENGINES_PATH 环境变量控制。

当此变量填充了文件系统上的有效路径时，TFTRT 转换器会将转换器构建的所有 TRT 引擎导出到此位置。这可以作为标准 TRT 引擎进行调查。

6.6. 阻止转换 Ops 以进行调试

当尝试识别 TFTRT 中的一个或多个有问题 Ops 时，阻止 Ops 转换为 TRT 是一种有用的调试工具。这由环境变量 TF_TRT_OP_DENYLIST 控制。这可以采用许多值，并且必须是完全字符串匹配。例如

            

            
TF_TRT_OP_DENYLIST=Sub,Exp,Conv2D
        

将阻止 Sub、Exp 和 Conv2D 从 b 转换为 TRT。

6.7. 使用实验性功能

环境变量 TF_TRT_EXPERIMENTAL_FEATURES 可用于在 TF-TRT 中启用特定的实验性功能。

用法：TF_TRT_EXPERIMENTAL_FEATURES=feature_1,feature_2 TF-TRT 当前支持以下实验性功能

• disable_graph_freezing
• reject_all_fp_cast_ops

6.8. 覆盖 NMS 插件的 top_k 阈值

TRT NMS 插件 允许 top_k<=4096，其中 top_k = max(num_boxes, max_total_size)。用户可以通过设置 TF_TRT_ALLOW_NMS_TOPK_OVERRIDE=1 环境变量来覆盖此值，但这可能会导致精度损失。此变量在 TRT 8.2.1 之前有效，较新的 TRT 版本对 top_k 没有限制。

用法：TF_TRT_ALLOW_NMS_TOPK_OVERRIDE=1

6.9. 启用张量布局优化器

环境变量 TF_TRT_ENABLE_LAYOUT_OPTIMIZER 用于启用/禁用布局转换为 NCHW。默认情况下，布局优化器处于打开状态。

用法：TF_TRT_ENABLE_LAYOUT_OPTIMIZER=0

6.10. 允许回退到 TF 原生段执行

环境变量 TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION 允许在 TRTEngine 执行失败时继续执行图。

用法：TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION=1

6.11. 控制生成的引擎数量

环境变量 TF_TRT_MAX_ALLOWED_ENGINES 允许用户指定将在特定模型中生成的 TRT 引擎的最大数量。这对于分段较多的网络非常方便，并且可以防止生成数十个引擎，从而可能提高性能。

用法

• 不限制引擎数量：TF_TRT_ENABLE_LAYOUT_OPTIMIZER=0
• 限制创建的引擎数量：TF_TRT_MAX_ALLOWED_ENGINES=<N> 其中 <N> 是创建的引擎的最大数量。如果图中存在超过 N 个段，则最大的 N 个段将被转换并创建引擎。

6.12. 可视化 TF-TRT 图

环境变量 TF_TRT_EXPORT_GRAPH_VIZ_PATH 启用并控制将在何处生成网络的图形可视化。

这是一个用于可视化 dot 格式转换后的 TRT 图表的实用工具。可以使用各种工具来查看 dot 文件。

用法: TF_TRT_EXPORT_GRAPH_VIZ_PATH=/path/to/graphviz-output

例如，下图显示了我们之前描述的 mnist 示例的可视化

TensorFlow-TensorRT GitHub 存储库是提交问题的最佳场所。在 GitHub 中发布问题时，请遵循 Stack Overflow 文档中概述的流程。确保发布的示例是

• 最小化 – 使用尽可能少的代码，同时仍然重现相同的问题。
• 完整 – 提供重现问题所需的所有部分。检查是否可以去除外部依赖项并仍然显示问题。我们在重现问题上花费的时间越少，我们就越有时间修复它。
• 可验证 – 测试您即将提供的代码，以确保它重现了问题。删除所有其他与您的请求/问题无关的问题。

8.1. 内存管理

TensorRT 将权重和激活存储在 GPU 上。存储在 TRTEngineOp 缓存中的每个引擎的大小与权重的大小（原始模型大小）大致相同。TensorRT 通过 TensorFlow 分配器分配内存，因此，所有 TensorFlow 内存配置限制 也适用于 TensorRT。

内存使用量高度依赖于模型，因为模型可能很大（例如 transformers）或很小（简单的图像分类模型.如果您观察到您的推理或引擎构建内存不足，请尝试减少批大小。

8.2. 最大缓存引擎数

这是一个主要用于启用 dynamic_op_mode 的隐式批处理模式的参数，不适用于动态形状模式。这指定了可以为动态 TRT 操作 (dynamic_op_mode) 缓存的最大 TRT 引擎数量。为动态维度创建的 TRT 引擎将被缓存。如果缓存的引擎数量已达到最大值，但没有一个引擎支持输入形状，则 TRTEngineOp 将回退以运行与 TRTEngineOp 对应的原始 TF 子图（原生段）。

TensorRT 引擎缓存在位于 TRTEngineOp op 中的 LRU 缓存中。如果缓存为空，或者缓存中不存在给定输入形状的引擎，则会创建一个新引擎。您可以使用参数 maximum_cached_engines 控制缓存的引擎数量。TensorRT 使用输入批大小作为选择性能最高的 CUDA 内核的参数之一。如果满足以下条件，则引擎可以重复用于新输入，

• 引擎批大小大于或等于新输入的批大小，并且
• 非批处理维度与新输入匹配。

如果您想保守地使用内存，请将 maximum_cached_engines 设置为 1，以强制在每次创建新引擎时都逐出任何现有缓存。另一方面，如果您的主要目标是减少延迟，则增加 maximum_cached_engines 以尽可能防止重新创建引擎。缓存更多引擎会占用机器上更多的资源，但是，对于典型模型来说，这不是问题。

注意

本文档仅供参考，不得视为对产品的特定功能、条件或质量的保证。NVIDIA Corporation（“NVIDIA”）对本文档中包含的信息的准确性或完整性不作任何明示或暗示的陈述或保证，并且对本文档中包含的任何错误不承担任何责任。NVIDIA 对使用此类信息或因使用此类信息而可能导致的侵犯第三方专利或其他权利的后果或使用不承担任何责任。本文档不承诺开发、发布或交付任何材料（下文定义）、代码或功能。

NVIDIA 保留随时修改、增强、改进和对本文档进行任何其他更改的权利，恕不另行通知。

客户在下订单前应获取最新的相关信息，并应验证此类信息是否为最新且完整。

NVIDIA 产品的销售受订单确认时提供的 NVIDIA 标准销售条款和条件约束，除非 NVIDIA 和客户的授权代表签署的个别销售协议（“销售条款”）另有约定。NVIDIA 特此明确反对将任何客户一般条款和条件应用于购买本文档中引用的 NVIDIA 产品。本文档未直接或间接形成任何合同义务。

NVIDIA 产品并非设计、授权或保证适用于医疗、军事、航空、航天或生命支持设备，也不适用于 NVIDIA 产品发生故障或失灵可能合理预期会导致人身伤害、死亡或财产或环境损害的应用。NVIDIA 对在上述设备或应用中包含和/或使用 NVIDIA 产品不承担任何责任，因此，此类包含和/或使用由客户自行承担风险。

NVIDIA 不作任何陈述或保证，保证基于本文档的产品适用于任何特定用途。NVIDIA 不一定对每种产品的所有参数进行测试。客户全权负责评估和确定本文档中包含的任何信息的适用性，确保产品适合客户计划的应用和用途，并为该应用执行必要的测试，以避免应用或产品出现缺陷。客户产品设计中的缺陷可能会影响 NVIDIA 产品的质量和可靠性，并可能导致超出本文档中包含的附加或不同条件和/或要求。NVIDIA 对以下原因可能导致或归因于的任何缺陷、损害、成本或问题不承担任何责任：(i) 以任何违反本文档的方式使用 NVIDIA 产品或 (ii) 客户产品设计。

本文档未授予 NVIDIA 专利权、版权或 NVIDIA 其他知识产权项下的任何明示或暗示许可。NVIDIA 发布的有关第三方产品或服务的信息不构成 NVIDIA 授予使用此类产品或服务的许可，也不构成对其的保证或认可。使用此类信息可能需要获得第三方在其专利或其他知识产权项下的许可，或获得 NVIDIA 在其专利或其他知识产权项下的许可。

只有在事先获得 NVIDIA 书面批准的情况下，才可以复制本文档中的信息，复制时不得进行更改，并且必须完全遵守所有适用的出口法律法规，并且必须附带所有相关的条件、限制和注意事项。

本文档以及所有 NVIDIA 设计规范、参考板、文件、图纸、诊断程序、列表和其他文档（统称为“材料”）均按“原样”提供。NVIDIA 对材料不作任何明示、暗示、法定或其他形式的保证，并且明确声明不承担所有关于不侵权、适销性和针对特定用途的适用性的暗示保证。在法律未禁止的范围内，在任何情况下，NVIDIA 均不对因使用本文档而引起的任何损害（包括但不限于任何直接、间接、特殊、附带、惩罚性或后果性损害，无论如何造成以及无论责任理论如何）承担责任，即使 NVIDIA 已被告知可能发生此类损害。尽管客户可能因任何原因遭受任何损害，但 NVIDIA 对本文所述产品的客户承担的累计总责任应根据产品的销售条款进行限制。

HDMI

HDMI、HDMI 标志和 High-Definition Multimedia Interface 是 HDMI Licensing LLC 的商标或注册商标。

OpenCL

OpenCL 是 Apple Inc. 的商标，已获得 Khronos Group Inc. 的许可使用。

商标

NVIDIA、NVIDIA 徽标以及 cuBLAS、CUDA、DALI、DGX、DGX-1、DGX-2、DGX Station、DLProf、Jetson、Kepler、Maxwell、NCCL、Nsight Compute、Nsight Systems、NvCaffe、PerfWorks、Pascal、SDK Manager、Tegra、TensorRT、Triton Inference Server、Tesla、TF-TRT 和 Volta 是 NVIDIA Corporation 在美国和其他国家/地区的商标和/或注册商标。其他公司和产品名称可能是与其相关的各自公司的商标。