NVIDIA TensorRT与Apache Beam SDK的集成

为大规模运行机器学习模型而加载和预处理数据通常需要将数据处理框架和推理机无缝拼接在一起。

在这篇文章中，我们将介绍 NVIDIA TensorRT 与 Apache Beam SDK 的集成，并展示如何将复杂的推理场景完全封装在数据处理管道中。我们还演示了如何通过几行代码处理来自批处理和流传输源的 TB 数据，以实现高吞吐量和低延迟模型推断。

NVIDIA TensorRT 是一个促进高性能机器学习推理的 SDK 。它设计用于 TensorFlow 、 PyTorch 和 MXNet 等深度学习框架。它专门致力于优化和运行一个经过训练的神经网络，以便在 NVIDIA 上高效地进行推理 GPU 。 TensorRT 可以通过多种优化最大化推理吞吐量，同时保持模型精度，包括模型量化、层和张量融合、内核自动调整、多流执行和有效的张量内存使用。

Dataflow 是一个无操作、无服务器的数据处理平台，经过 15 年以上的生产实践证明，可以批量或实时处理数据，用于分析、 ML 和应用程序用例。这些通常包括将预训练的模型合并到数据管道中。无论使用情况如何， Apache Beam 作为其 SDK 的使用使 DataFlow 能够利用强大的社区，简化您的数据架构，并通过 ML 提供见解。

构建 TensorRT 推理引擎

要将 TensorRT 与 Apache Beam 一起使用，在此阶段，您需要 converted TensorRT engine file from a trained model. 以下是如何将 TensorFlow 对象检测 SSD MobileNet v2 320 × 320 模型转换为 ONNX ，从 ONNX 构建 TensorRT 引擎，并在本地运行引擎。

将 TF 模型转换为 ONNX

要将 TensorFlow 对象检测 SSD MobileNet v2 320 × 320 转换为 ONNX ，请使用 one of the TensorRT example converters 。如果系统具有将在数据流中用于推断的相同 GPU ，则可以在本地系统上执行此操作。

要准备环境，请按照 Setup 中的说明进行操作。本文遵循本指南，直到 Create ONNX Graph 。使用– batch _ size 1 作为示例，我们将介绍仅使用批处理大小 1 的进一步工作。您可以将最终– onnx 文件命名为 ssd _ mobilenet _ v2 _ 320×320 _ coco17 _ tpu-8.onnx 。构建和运行在 GCP 中处理。

确保您设置的 GCP 项目具有正确的凭据和对 Dataflow 、 Google 云存储（ GCS ）和 Google 计算引擎（ GCE ）的 API 访问权限。有关详细信息，请参见 Create a Dataflow pipeline using Python 。

启动 GCE VM

您需要一台包含以下已安装资源的计算机：

NVIDIA T4 Tensor 核心 GPU

GPU 驱动器

Docker 公司

NVIDIA 容器工具包

您可以通过 creating a new GCE VM 执行此操作。按照说明操作，但使用以下设置：

Name: tensorrt-demo

GPU type: 环境 T4

Number of GPUs: 1

Machine type: n1-standard-2

如果你知道你使用的是大型模型，你可能需要一台更强大的机器。

在 Boot disk 部分，选择 CHANGE ，然后转到 PUBLIC IMAGES 选项卡。对于 Operating system ，选择 Linux 上的深度学习 。有很多版本，但请确保选择 CUDA 版本。版本 基于 Debian 10 的深度学习 VM 与 M98 适用于此示例。

其他设置可以保留为默认值。

接下来， connect to the VM using SSH 。如果系统提示您安装 NVIDIA 驱动程序，请安装。

在 VM 内部，运行以下命令以创建一些稍后使用的目录：

 

mkdir models
mkdir tensorrt_engines

 

有关详细信息，请参见 Create a VM with attached GPUs 。

建立形象

您需要一个自定义容器，其中包含执行 TensorRT 代码所需的依赖项： CUDA 、 cuDNN 和 TensorRT 。

您可以将以下示例 Dockerfile 复制到新文件中，并将其命名为tensor_rt.dockerfile .

 

ARG BUILD_IMAGE=nvcr.io/nvidia/tensorrt:22.09-py3

FROM ${BUILD_IMAGE} 

ENV PATH="/usr/src/tensorrt/bin:${PATH}"

WORKDIR /workspace

RUN pip install --no-cache-dir apache-beam[gcp]==2.42.0
COPY --from=apache/beam_python3.8_sdk:2.42.0 /opt/apache/beam /opt/apache/beam

RUN pip install --upgrade pip 
    && pip install torch>=1.7.1 
    && pip install torchvision>=0.8.2 
    && pip install pillow>=8.0.0 
    && pip install transformers>=4.18.0 
    && pip install cuda-python

ENTRYPOINT [ "/opt/apache/beam/boot" ]

 

查看 Docker file used for testing in the Apache Beam repo 。请记住，可能会有比本文所用版本更高版本的 Beam 可用。

通过在本地或在 GCE VM 中运行以下命令来构建映像：

 

docker build -f tensor_rt.dockerfile -t tensor_rt .

 

如果在本地执行此操作，请执行以下步骤。否则，您可以跳到下一节。

只有当您在不同于您打算构建 TensorRT 引擎的机器上创建图像时，才需要以下命令。对于这篇文章，请使用 Google Container Registry 。将图像标记为用于项目的 URI ，然后推送到注册表。确保用适当的值替换GCP_PROJECT和MY_DIR。

 

docker tag tensor_rt us.gcr.io/{GCP_PROJECT}/{MY_DIR}/tensor_rt
docker push us.gcr.io/{GCP_PROJECT}/{MY_DIR}/tensor_rt

 

创建 TensorRT 引擎

只有当您在不同于要构建 TensorRT 引擎的机器上创建图像时，才需要以下命令。从注册表中提取 TensorRT 图像：

 

docker pull us.gcr.io/{GCP_PROJECT}/{MY_DIR}/tensor_rt
docker tag us.gcr.io/{GCP_PROJECT}/{MY_DIR}/tensor_rt tensor_rt

 

如果 ONNX 模型不在 GCE VM 中，您可以将其从本地计算机复制到/models 目录：

 

gcloud compute scp ~/Downloads/ssd_mobilenet_v2_320x320_coco17_tpu-8.onnx tensorrt-demo:~/models --zone=us-central1-a

 

现在，您应该在 VM 中拥有 ONNX 模型和构建的 Docker 映像。现在是时候同时使用它们了。

以交互方式启动 Docker 容器：

 

docker run --rm -it --gpus all -v /home/{username}/:/mnt tensor_rt bash

 

从 ONNX 文件创建 TensorRT 引擎：

 

trtexec --onnx=/mnt/models/ssd_mobilenet_v2_320x320_coco17_tpu-8.onnx --saveEngine=/mnt/tensorrt_engines/ssd_mobilenet_v2_320x320_coco17_tpu-8.trt --useCudaGraph --verbose

 

现在，您应该可以在 VM 的/tensorrt_engines目录中看到ssd_mobilenet_v2_320x320_coco17_tpu-8.trt文件。

将 TensorRT 引擎上传至 GCS

将文件复制到 GCP 。如果您在将文件直接从 GCE 上传到 GCS 时遇到gsutil 问题，您可能必须首先将其复制到本地计算机。

 

gcloud compute scp tensorrt-demo:~/tensorrt_engines/ssd_mobilenet_v2_320x320_coco17_tpu-8.trt ~/Downloads/ --zone=us-central1-a

 

在 GCP 控制台中，将 TensorRT 引擎文件上传到您选择的 GCS 存储桶：

 

gs://{GCS_BUCKET}/ssd_mobilenet_v2_320x320_coco17_tpu-8.trt

 

本地测试 TensorRT 引擎

确保您拥有使用 TensorRT RunInference 的 Beam 管道。一个示例是 tensorrt_object_detection.py ，您可以通过在 GCE VM 中运行以下命令来遵循该示例。首先键入 Ctrl + D 退出 Docker 容器。

 

git clone https://github.com/apache/beam.git
cd beam/sdks/python
pip install --upgrade pip setuptools
pip install -r build-requirements.txt
pip install --user -e ."[gcp,test]"

 

您还创建了一个名为image_file_names.txt的文件，其中 包含图像的路径。图像可以在 GCS 之类的对象存储中，也可以在 GCE VM 中。

 

gs://{GCS_BUCKET}/000000289594.jpg
gs://{GCS_BUCKET}/000000000139.jpg

 

然后，运行以下命令：

 

docker run --rm -it --gpus all -v /home/{username}/:/mnt -w /mnt/beam/sdks/python tensor_rt python -m apache_beam.examples.inference.tensorrt_object_detection --input gs://{GCS_BUCKET}/tensorrt_image_file_names.txt --output /mnt/tensorrt_predictions.csv --engine_path gs://{GCS_BUCKET}/ssd_mobilenet_v2_320x320_coco17_tpu-8.trt

 

现在您应该看到一个名为tensorrt_predictions.csv . 的文件。每行都有用分号分隔的数据。

第一项是文件名。

第二项是字典列表，其中每个字典对应一个检测。

检测包含框坐标（ ymin 、 xmin 、 ymax 、 xmax ）、分数和类别。

有关如何在本地设置和运行 TensorRT RunInference 的更多信息，请遵循 Object Detection 部分中的说明。

TensorRT Support Guide 概述了 GitHub 和产品包中所有支持的 NVIDIA TensorRT 8.5.1 示例。这些示例旨在展示如何在众多用例中使用 TensorRT ，同时突出显示界面的不同功能。这些示例在推荐器、机器理解、字符识别、图像分类和对象检测等用例中特别有用。

使用 DataFlow RunInference 运行 TensorRT 引擎

现在您有了 TensorRT 引擎，就可以在 Dataflow 上运行管道了。

下面的代码示例是管道的一部分，您可以使用TensorRTEngineHandlerNumPy加载 TensorRT 引擎并设置其他推断参数。然后读取图像，进行预处理以将关键点附加到图像，进行预测，然后写入 GCS 中的文件。

有关完整代码示例的更多信息，请参见 tensorrt_object_detection.py 。

 

  engine_handler = KeyedModelHandler(
      TensorRTEngineHandlerNumPy(
          min_batch_size=1,
          max_batch_size=1,
          engine_path=known_args.engine_path))

  with beam.Pipeline(options=pipeline_options) as p:
    filename_value_pair = (
        p
        | 'ReadImageNames' >> beam.io.ReadFromText(known_args.input)
        | 'ReadImageData' >> beam.Map(
            lambda image_name: read_image(
                image_file_name=image_name, path_to_dir=known_args.images_dir))
        | 'AttachImageSizeToKey' >> beam.Map(attach_im_size_to_key)
        | 'PreprocessImages' >> beam.MapTuple(
            lambda file_name, data: (file_name, preprocess_image(data))))
    predictions = (
        filename_value_pair
        | 'TensorRTRunInference' >> RunInference(engine_handler)
        | 'ProcessOutput' >> beam.ParDo(PostProcessor()))

    _ = (
        predictions | "WriteOutputToGCS" >> beam.io.WriteToText(
            known_args.output,
            shard_name_template='',
            append_trailing_newlines=True))

 

确保您已完成上一节中提到的 Google Cloud 设置。您还必须具有 Beam SDK installed 。

要在 Dataflow 上运行此作业，请在本地运行以下命令：

 

python -m apache_beam.examples.inference.tensorrt_object_detection 
--input gs://{GCP_PROJECT}/image_file_names.txt 
--output gs://{GCP_PROJECT}/predictions.txt 
--engine_path gs://{GCP_PROJECT}/ssd_mobilenet_v2_320x320_coco17_tpu-8.trt 
--runner DataflowRunner 
--experiment=use_runner_v2 
--machine_type=n1-standard-4 
--experiment="worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver" 
--disk_size_gb=75 
--project {GCP_PROJECT} 
--region us-central1 
--temp_location gs://{GCP_PROJECT}/tmp/ 
--job_name tensorrt-object-detection 
--sdk_container_image="us.gcr.io/{GCP_PROJECT}/{MY_DIR}/tensor_rt tensor_rt"

 

根据模型的大小限制，您可能需要调整 machine _ type 、 GPU 的类型和计数或 disk _ size _ gb 。有关梁管道选项的详细信息，请参见 Set Dataflow pipeline options 。

TensorRT 和 TensorFlow 目标检测基准

为了进行基准测试，我们决定在前面提到的 SSD MobileNet v2 320 × 320 模型的 TensorRT 和 TensorFlow 对象检测版本之间进行比较。

在 TensorRT 和 TensorFlow 对象检测版本中，每个推理调用都是定时的。我们计算了平均 5000 个推断调用，由于延迟增加，没有考虑前 10 个图像。我们使用的 SSD 型号是小型型号。当您的模型可以充分利用 GPU 时，您将观察到更好的加速。

首先，我们将 TensorFlow 和 TensorRT 之间的直接性能加速与本地基准进行了比较。我们旨在证明 TensorRT 上降低精度模式的额外优势。

Framework and precision	Inference latency (ms)
TensorFlow Object Detection FP32 (end-to-end)	29.47 ms
TensorRT FP32 (end-to-end)	3.72 ms
TensorRT FP32 (GPU compute)	2.39 ms
TensorRT FP16 (GPU compute)	1.48 ms
TensorRT INT8 (GPU compute)	1.34 ms

表 1 。 TensorRT 上的直接性能加速

TensorRT FP32 的总体加速为 7.9x 。端到端包括数据副本，而 GPU 计算仅包括实际推断时间。我们这样做是因为示例模型很小。在这种情况下，端到端 TensorRT 延迟主要是数据拷贝。在更大的模型中使用不同的精度可以看到更显著的端到端性能改进，尤其是在推理计算是瓶颈而不是数据拷贝的情况下。

FP16 比 FP32 快 1.6 倍，没有精度损失。 INT8 比 FP32 快 1.8 倍，但有时精度会降低，需要校准过程。精度下降是特定于模型的，因此尝试您的精度并查看产生的精度总是很好的。

使用 NVIDIA QAT 工具包的量化网络也可以缓解此问题。有关详细信息，请参见 Accelerating Quantized Networks with the NVIDIA QAT Toolkit for TensorFlow and NVIDIA TensorRT 和 NVIDIA TensorRT Developer Guide .

数据流基准测试

在 Dataflow 中，使用早期实验中生成的 TensorRT 引擎，我们使用以下配置运行：n1-standard-4 machine、disk_size_gb=75和 10 个工作人员。

为了模拟通过PubSub进入 Dataflow 的数据流，我们将批大小设置为 1 。这是通过将ModelHandlers设置为最小和最大批量大小为 1 来实现的。

	Stage with RunInference	Mean inference_batch_latency_micro_secs
TensorFlow with T4 GPU	12 min 43 sec	99,242
TensorRT with T4 GPU	7 min 20 sec	10,836

表 2 。数据流基准

Dataflow runner 将管道分解为多个阶段。通过查看包含推理调用的阶段，而不是读取和写入数据的其他阶段，可以更好地了解RunInference的性能。这在 Stage with RunInference 列中。

对于这个度量， TensorRT 只花费 TensorFlow 运行时间的 57% 。如果你适应一个完全使用 GPU 处理能力的更大模型，你预计加速度会增长。

度量推理_ batch _ latency _ micro _ secs 是对一批示例执行推理所需的时间（以微秒为单位），即调用model_handler.run_inference的时间。这随着时间的推移而变化，这取决于BatchElements的动态批处理决策以及元素的特定值或dtype值。对于这个度量，您可以看到 TensorRT 比 TensorFlow 快 9.2 倍。

结论

在这篇文章中，我们演示了如何通过无缝拼接数据处理框架（ Apache Beam ）和推理引擎（ TensorRT ）来大规模运行机器学习模型。我们提供了一个端到端的示例，说明如何将推理工作负载完全集成到数据处理管道中。

这种集成实现了一种新的推理流水线，该流水线有助于通过更好的 NVIDIA GPU 利用率和大大提高的推理延迟和吞吐量来降低生产推理成本。使用许多现成的 TensorRT 样本，相同的方法可以应用于许多其他推断工作负载。未来，我们计划进一步自动化 TensorRT 引擎构建，并致力于 TensorRT 与 Apache Beam 的深度集成。