為大規模運行機器學習模型而加載和預處理數據通常需要將數據處理框架和推理機無縫拼接在一起。

在這篇文章中，我們將介紹 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 上降低精度模式的額外優勢。

表 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 來實現的。

表 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 的深度集成。