Jonathan Dekhtiar（NVIDIA）、Bixia Zheng（Google）、Shashank Verma（NVIDIA）、Chetan Tekur（NVIDIA）发布
TensorFlow-TensorRT (TF-TRT) 是 TensorFlow 和 TensorRT 的集成，它利用 TensorFlow 生态系统中的 NVIDIA GPU 推理优化。它提供了一个简单的 API，可以轻松地在 NVIDIA GPU 上实现显著的性能提升。该集成允许利用 TensorRT 中可能的优化，同时在遇到 TensorRT 不支持的模型部分时提供回退到原生 TensorFlow 的机制。
在我们之前关于 TF-TRT 集成的博客中，我们介绍了 TensorFlow 1.13 及更早版本的工作流程。这篇博客将介绍 TensorFlow 2.x 中的 TensorRT 集成，并演示使用最新 API 的示例工作流程。即使您是 TensorRT 集成的初学者，这篇博客也包含了您入门所需的所有信息。使用 TensorRT 集成已证明可以将性能提高 2.4 倍，而原生 TensorFlow 推理在 Nvidia T4 GPU 上的性能则可以提高 2.4 倍。
当启用 TF-TRT 时，在第一步中，会解析训练后的模型以将图分割为 TensorRT 支持 的子图和不支持的子图。然后，每个 TensorRT 支持的子图都封装在一个单独的特殊 TensorFlow 操作（TRTEngineOp）中。在第二步中，对于每个 TRTEngineOp 节点，都会构建一个优化的 TensorRT 引擎。TensorRT 不支持的子图保持不变，并由 TensorFlow 运行时处理。图 1 说明了这一点。
TF-TRT 允许利用 TensorFlow 的灵活性，同时利用可以应用于 TensorRT 支持的子图的优化。只有图的一部分经过优化并使用 TensorRT 执行，而 TensorFlow 执行剩余的图。
在图 1 所示的推理示例中，TensorFlow 执行 Reshape 操作和 Cast 操作。然后 TensorFlow 将 TRTEngineOp_0（预构建的 TensorRT 引擎）的执行传递给 TensorRT 运行时。
 |
| 图 1：TF-TRT 中图分割和构建 TRT 引擎的示例 |
在本节中，我们将通过示例了解典型的 TF-TRT 工作流程。
 |
| 图 2：仅在 TensorFlow 中执行推理时的工作流程图，以及使用已转换的 SavedModel 在 TensorFlow-TensorRT 中执行推理时的工作流程图 |
图 2 显示了原生 TensorFlow 中的标准推理工作流程，并将其与 TF-TRT 工作流程进行了对比。SavedModel 格式包含共享或部署训练模型所需的所有信息。在原生 TensorFlow 中，工作流程通常包括加载保存的模型，并使用 TensorFlow 运行时执行推理。在 TF-TRT 中，会涉及一些额外的步骤，包括对模型的 TensorRT 支持的子图应用 TensorRT 优化，以及可选地预构建 TensorRT 引擎。
首先，我们创建一个对象来保存转换参数，包括精度模式。精度模式用于指示 TF-TRT 可以用来实现 TensorFlow 操作的最低精度（例如 FP32、FP16 或 INT8）。然后，我们创建一个转换器对象，它接受转换参数和来自保存模型的输入。请注意，在 TensorFlow 2.x 中，TF-TRT 仅支持以 TensorFlow SavedModel 格式保存的模型。
接下来，当我们调用转换器 convert() 方法时，TF-TRT 会通过用 TRTEngineOp 替换图中的 TensorRT 兼容部分来转换图。为了在运行时获得更好的性能，转换器 build() 方法可用于提前创建 TensorRT 执行引擎。build() 方法要求在构建优化后的 TensorRT 执行引擎之前知道输入数据的形状。如果不知道输入数据的形状，则可以在运行时根据输入数据可用时构建 TensorRT 执行引擎。TensorRT 执行引擎应该在与执行推理的设备类型相同的 GPU 上构建，因为构建过程是特定于 GPU 的。例如，为 Nvidia A100 GPU 构建的执行引擎不能在 Nvidia T4 GPU 上运行。
最后，通过调用 save 方法，可以将 TF-TRT 转换后的模型保存到磁盘。与本节中提到的工作流程步骤相对应的代码显示在下面的代码块中
from tensorflow.python.compiler.tensorrt import trt_convert as trt
# Conversion Parameters
conversion_params = trt.TrtConversionParams(
precision_mode=trt.TrtPrecisionMode.<FP32 or FP16>)
converter = trt.TrtGraphConverterV2(
input_saved_model_dir=input_saved_model_dir,
conversion_params=conversion_params)
# Converter method used to partition and optimize TensorRT compatible segments
converter.convert()
# Optionally, build TensorRT engines before deployment to save time at runtime
# Note that this is GPU specific, and as a rule of thumb, we recommend building at runtime
converter.build(input_fn=my_input_fn)
# Save the model to the disk
converter.save(output_saved_model_dir)从上面的代码示例可以看出，build() 方法需要一个与输入数据形状相对应的输入函数。下面是一个输入函数的示例
# input_fn: a generator function that yields input data as a list or tuple,
# which will be used to execute the converted signature to generate TensorRT
# engines. Example:
def my_input_fn():
# Let's assume a network with 2 input tensors. We generate 3 sets
# of dummy input data:
input_shapes = [[(1, 16), (2, 16)], # min and max range for 1st input list
[(2, 32), (4, 32)], # min and max range for 2nd list of two tensors
[(4, 32), (8, 32)]] # 3rd input list
for shapes in input_shapes:
# return a list of input tensors
yield [np.zeros(x).astype(np.float32) for x in shapes]与 FP32 和 FP16 相比，INT8 需要额外的校准数据来确定最佳量化阈值。当转换参数中的精度模式为 INT8 时，我们需要在 convert() 方法调用中提供一个输入函数。此输入函数类似于提供给 build() 方法的输入函数。此外，通过传递给 convert() 方法的输入函数生成的校准数据应生成与推理期间看到的实际数据在统计上相似的數據。
from tensorflow.python.compiler.tensorrt import trt_convert as trt
conversion_params = trt.TrtConversionParams(
precision_mode=trt.TrtPrecisionMode.INT8)
converter = trt.TrtGraphConverterV2(
input_saved_model_dir=input_saved_model_dir,
conversion_params=conversion_params)
# requires some data for calibration
converter.convert(calibration_input_fn=my_input_fn)
# Optionally build TensorRT engines before deployment.
# Note that this is GPU specific, and as a rule of thumb we recommend building at runtime
converter.build(input_fn=my_input_fn)
converter.save(output_saved_model_dir)这篇博客的其余部分将展示获取 TensorFlow 2.x ResNet-50 模型、训练它、保存它、使用 TF-TRT 优化它，最后将其部署以进行推理的工作流程。我们还将比较使用 TensorFlow 原生模式和 TF-TRT 在三种精度模式（FP32、FP16 和 INT8）下执行推理的吞吐量。
首先，需要从 TensorFlow github 存储库下载最新版本的 ResNet-50 模型
# Adding the git remote and fetch the existing branches
$ git clone --depth 1 https://github.com/tensorflow/models.git .
# List the files and directories present in our working directory
$ ls -al
rwxrwxr-x user user 4 KiB Wed Sep 30 15:31:05 2020 ./
rwxrwxr-x user user 4 KiB Wed Sep 30 15:30:45 2020 ../
rw-rw-r-- user user 337 B Wed Sep 30 15:31:05 2020 AUTHORS
rw-rw-r-- user user 1015 B Wed Sep 30 15:31:05 2020 CODEOWNERS
rwxrwxr-x user user 4 KiB Wed Sep 30 15:31:05 2020 community/
rw-rw-r-- user user 390 B Wed Sep 30 15:31:05 2020 CONTRIBUTING.md
rwxrwxr-x user user 4 KiB Wed Sep 30 15:31:15 2020 .git/
rwxrwxr-x user user 4 KiB Wed Sep 30 15:31:05 2020 .github/
rw-rw-r-- user user 1 KiB Wed Sep 30 15:31:05 2020 .gitignore
rw-rw-r-- user user 1 KiB Wed Sep 30 15:31:05 2020 ISSUES.md
rw-rw-r-- user user 11 KiB Wed Sep 30 15:31:05 2020 LICENSE
rwxrwxr-x user user 4 KiB Wed Sep 30 15:31:05 2020 official/
rwxrwxr-x user user 4 KiB Wed Sep 30 15:31:05 2020 orbit/
rw-rw-r-- user user 3 KiB Wed Sep 30 15:31:05 2020 README.md
rwxrwxr-x user user 4 KiB Wed Sep 30 15:31:06 2020 research/如前几节所述，在本示例中，我们将使用 Docker 存储库中可用的最新 TensorFlow 容器。用户不需要任何额外的安装步骤，因为 TensorRT 集成已包含在容器中。拉取容器并启动容器的步骤如下所示
$ docker pull tensorflow/tensorflow:latest-gpu
# Please ensure that the Nvidia Container Toolkit is installed before running the following command
$ docker run -it --rm \
--gpus="all" \
--shm-size=2g --ulimit memlock=-1 --ulimit stack=67108864 \
--workdir /workspace/ \
-v "$(pwd):/workspace/" \
-v "</path/to/save/data/>:/data/" \ # This is the path that will hold the training data
tensorflow/tensorflow:latest-gpu在容器内部，我们可以验证我们是否可以访问相关文件以及我们要针对的 Nvidia GPU
# Let's first test that we can access the ResNet-50 code that we previously downloaded
$ ls -al
drwxrwxr-x 8 1000 1000 4096 Sep 30 22:31 .git
drwxrwxr-x 3 1000 1000 4096 Sep 30 22:31 .github
-rw-rw-r-- 1 1000 1000 1104 Sep 30 22:31 .gitignore
-rw-rw-r-- 1 1000 1000 337 Sep 30 22:31 AUTHORS
-rw-rw-r-- 1 1000 1000 1015 Sep 30 22:31 CODEOWNERS
-rw-rw-r-- 1 1000 1000 390 Sep 30 22:31 CONTRIBUTING.md
-rw-rw-r-- 1 1000 1000 1115 Sep 30 22:31 ISSUES.md
-rw-rw-r-- 1 1000 1000 11405 Sep 30 22:31 LICENSE
-rw-rw-r-- 1 1000 1000 3668 Sep 30 22:31 README.md
drwxrwxr-x 2 1000 1000 4096 Sep 30 22:31 community
drwxrwxr-x 12 1000 1000 4096 Sep 30 22:31 official
drwxrwxr-x 3 1000 1000 4096 Sep 30 22:31 orbit
drwxrwxr-x 23 1000 1000 4096 Sep 30 22:31 research
# Let's verify we can see our GPUs:
$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.XX.XX Driver Version: 450.XX.XX CUDA Version: 11.X |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Tesla T4 On | 00000000:1A:00.0 Off | Off |
| 38% 52C P8 14W / 70W | 1MiB / 16127MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
现在我们可以开始训练 ResNet-50 了。为了避免花费数小时训练深度学习模型，本文将使用较小的 MNIST 数据集。但是，工作流程不会因使用 ImageNet 等更先进的数据集而改变。
# Install dependencies
$ pip install tensorflow_datasets tensorflow_model_optimization
# Download MNIST data and Train
$ python -m "official.vision.image_classification.mnist_main" \
--model_dir=./checkpoints \
--data_dir=/data \
--train_epochs=10 \
--distribution_strategy=one_device \
--num_gpus=1 \
--download
# Let’s verify that we have the trained model saved on our machine.
$ ls -al checkpoints/
-rw-r--r-- 1 root root 87 Sep 30 22:34 checkpoint
-rw-r--r-- 1 root root 6574829 Sep 30 22:34 model.ckpt-0001.data-00000-of-00001
-rw-r--r-- 1 root root 819 Sep 30 22:34 model.ckpt-0001.index
[...]
-rw-r--r-- 1 root root 6574829 Sep 30 22:34 model.ckpt-0010.data-00000-of-00001
-rw-r--r-- 1 root root 819 Sep 30 22:34 model.ckpt-0010.index
drwxr-xr-x 4 root root 4096 Sep 30 22:34 saved_model
drwxr-xr-x 3 root root 4096 Sep 30 22:34 train
drwxr-xr-x 2 root root 4096 Sep 30 22:34 validation
训练后，Google 的 ResNet-50 代码将模型以 SavedModel 格式导出到以下路径：checkpoints/saved_model/.
以下示例 代码 可用作将您自己的训练模型导出为 TensorFlow SavedModel 的参考。
import numpy as np
import tensorflow as tf
from tensorflow import keras
def get_model():
# Create a simple model.
inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = keras.Model(inputs, outputs)
model.compile(optimizer="adam", loss="mean_squared_error")
return model
model = get_model()
# Train the model.
test_input = np.random.random((128, 32))
test_target = np.random.random((128, 1))
model.fit(test_input, test_target)
# Calling `save('my_model')` creates a SavedModel folder `my_model`.
model.save("my_model")我们可以验证 Google 的 ResNet-50 脚本生成的 SavedModel 是否可读且正确
$ ls -al checkpoints/saved_model
drwxr-xr-x 2 root root 4096 Sep 30 22:49 assets
-rw-r--r-- 1 root root 118217 Sep 30 22:49 saved_model.pb
drwxr-xr-x 2 root root 4096 Sep 30 22:49 variables
$ saved_model_cli show --dir checkpoints/saved_model/ --tag_set serve --signature_def serving_default
MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:
The given SavedModel SignatureDef contains the following input(s):
inputs['input_1'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 28, 28, 1)
name: serving_default_input_1:0
The given SavedModel SignatureDef contains the following output(s):
outputs['dense_1'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 10)
name: StatefulPartitionedCall:0
Method name is: tensorflow/serving/predict现在我们已经验证了我们的 SavedModel 已正确保存，我们可以继续使用 TF-TRT 加载它以进行推理。
在本节中，我们将介绍在 NVIDIA GPU 上使用 TF-TRT 部署保存的 ResNet-50 模型的步骤。如前所述，我们首先使用 convert 方法将 SavedModel 转换为 TF-TRT 模型，然后加载模型。
# Convert the SavedModel
converter = trt.TrtGraphConverterV2(input_saved_model_dir=path)
converter.convert()
# Save the converted model
converter.save(converted_model_path)
# Load converted model and infer
model = tf.saved_model.load(converted_model_path)
func = root.signatures['serving_default']
output = func(input_tensor)为了简单起见，我们将使用脚本执行推理 (tf2_inference.py)。我们将从 github.com 下载脚本并将其放在与之前相同的 Docker 容器的工作目录“/workspace/”中。之后，我们可以执行脚本
$ wget https://raw.githubusercontent.com/tensorflow/tensorrt/master/tftrt/blog_posts/Leveraging%20TensorFlow-TensorRT%20integration%20for%20Low%20latency%20Inference/tf2_inference.py
$ ls
AUTHORS CONTRIBUTING.md LICENSE checkpoints data orbit tf2_inference.py
CODEOWNERS ISSUES.md README.md community official research
$ python tf2_inference.py --use_tftrt_model --precision fp16
=========================================
Inference using: TF-TRT …
Batch size: 512
Precision: fp16
=========================================
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
TrtConversionParams(rewriter_config_template=None, max_workspace_size_bytes=8589934592, precision_mode='FP16', minimum_segment_size=3, is_dynamic_op=True, maximum_cached_engines=100, use_calibration=True, max_batch_size=512, allow_build_at_runtime=True)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
Processing step: 0100 ...
Processing step: 0200 ...
[...]
Processing step: 9900 ...
Processing step: 10000 ...
Average step time: 2.1 msec
Average throughput: 244248 samples/sec类似地，我们可以为 INT8 和 FP32 运行推理
$ python tf2_inference.py --use_tftrt_model --precision int8
$ python tf2_inference.py --use_tftrt_model --precision fp32您还可以运行未经修改的 SavedModel，而无需任何 TF-TRT 加速。
$ python tf2_inference.py --use_native_tensorflow
=========================================
Inference using: Native TensorFlow …
Batch size: 512
=========================================
Processing step: 0100 ...
Processing step: 0200 ...
[...]
Processing step: 9900 ...
Processing step: 10000 ...
Average step time: 4.1 msec
Average throughput: 126328 samples/sec此运行是在 NVIDIA T4 GPU 上执行的。相同的工作流程将在任何 NVIDIA GPU 上运行。
对代码进行最小的更改以利用 TF-TRT 可以带来显著的性能提升。例如，使用本文中的推理脚本，在 NVIDIA T4 GPU 上使用 512 的批次大小，我们观察到 TF-TRT FP16 比原生 TensorFlow 快近 2 倍，而 TF-TRT INT8 比原生 TensorFlow 快 2.4 倍。获得的加速量可能因各种因素而异，例如所使用的模型、批次大小、数据集中的图像大小和格式以及任何 CPU 瓶颈。
总之，在本博客中，我们展示了 TF-TRT 提供的加速。此外，使用 TF-TRT，我们可以使用完整的 TensorFlow Python API 和交互式环境（如 Jupyter Notebook 或 Google Colab）。
TF-TRT 用户指南列出了 操作符，这些操作符在 TensorRT 兼容的子图中受支持。不在此列表中的操作符将由原生 TensorFlow 运行时执行。
我们鼓励您自己尝试一下，如果您遇到问题，请在这里 打开一个问题。
2021年1月28日 — 发布者：Jonathan Dekhtiar（NVIDIA），郑碧霞（Google），Shashank Verma（NVIDIA），Chetan Tekur（NVIDIA）
TensorFlow-TensorRT (TF-TRT) 是 TensorFlow 与 TensorRT 的集成，它利用 TensorFlow 生态系统中 NVIDIA GPU 的推理优化。它提供了一个简单的 API，只需付出最小的努力即可在 NVIDIA GPU 上获得显著的性能提升。这种集成允许…
