NVIDIA TensorRT與Apache Beam SDK的集成

為大規(guī)模運(yùn)行機(jī)器學(xué)習(xí)模型而加載和預(yù)處理數(shù)據(jù)通常需要將數(shù)據(jù)處理框架和推理機(jī)無縫拼接在一起。

在這篇文章中，我們將介紹 NVIDIA TensorRT 與 Apache Beam SDK 的集成，并展示如何將復(fù)雜的推理場景完全封裝在數(shù)據(jù)處理管道中。我們還演示了如何通過幾行代碼處理來自批處理和流傳輸源的 TB 數(shù)據(jù)，以實(shí)現(xiàn)高吞吐量和低延遲模型推斷。

NVIDIA TensorRT 是一個(gè)促進(jìn)高性能機(jī)器學(xué)習(xí)推理的 SDK 。它設(shè)計(jì)用于 TensorFlow 、 PyTorch 和 MXNet 等深度學(xué)習(xí)框架。它專門致力于優(yōu)化和運(yùn)行一個(gè)經(jīng)過訓(xùn)練的神經(jīng)網(wǎng)絡(luò)，以便在 NVIDIA 上高效地進(jìn)行推理 GPU 。 TensorRT 可以通過多種優(yōu)化最大化推理吞吐量，同時(shí)保持模型精度，包括模型量化、層和張量融合、內(nèi)核自動(dòng)調(diào)整、多流執(zhí)行和有效的張量內(nèi)存使用。

Dataflow 是一個(gè)無操作、無服務(wù)器的數(shù)據(jù)處理平臺(tái)，經(jīng)過 15 年以上的生產(chǎn)實(shí)踐證明，可以批量或?qū)崟r(shí)處理數(shù)據(jù)，用于分析、 ML 和應(yīng)用程序用例。這些通常包括將預(yù)訓(xùn)練的模型合并到數(shù)據(jù)管道中。無論使用情況如何， Apache Beam 作為其 SDK 的使用使 DataFlow 能夠利用強(qiáng)大的社區(qū)，簡化您的數(shù)據(jù)架構(gòu)，并通過 ML 提供見解。

構(gòu)建 TensorRT 推理引擎

要將 TensorRT 與 Apache Beam 一起使用，在此階段，您需要 converted TensorRT engine file from a trained model. 以下是如何將 TensorFlow 對象檢測 SSD MobileNet v2 320 × 320 模型轉(zhuǎn)換為 ONNX ，從 ONNX 構(gòu)建 TensorRT 引擎，并在本地運(yùn)行引擎。

將 TF 模型轉(zhuǎn)換為 ONNX

要將 TensorFlow 對象檢測 SSD MobileNet v2 320 × 320 轉(zhuǎn)換為 ONNX ，請使用 one of the TensorRT example converters 。如果系統(tǒng)具有將在數(shù)據(jù)流中用于推斷的相同 GPU ，則可以在本地系統(tǒng)上執(zhí)行此操作。

要準(zhǔn)備環(huán)境，請按照 Setup 中的說明進(jìn)行操作。本文遵循本指南，直到 Create ONNX Graph 。使用– batch _ size 1 作為示例，我們將介紹僅使用批處理大小 1 的進(jìn)一步工作。您可以將最終– onnx 文件命名為 ssd _ mobilenet _ v2 _ 320×320 _ coco17 _ tpu-8.onnx 。構(gòu)建和運(yùn)行在 GCP 中處理。

確保您設(shè)置的 GCP 項(xiàng)目具有正確的憑據(jù)和對 Dataflow 、 Google 云存儲(chǔ)（ GCS ）和 Google 計(jì)算引擎（ GCE ）的 API 訪問權(quán)限。有關(guān)詳細(xì)信息，請參見 Create a Dataflow pipeline using Python 。

啟動(dòng) GCE VM

您需要一臺(tái)包含以下已安裝資源的計(jì)算機(jī)：

NVIDIA T4 Tensor 核心 GPU

GPU 驅(qū)動(dòng)器

Docker 公司

NVIDIA 容器工具包

您可以通過 creating a new GCE VM 執(zhí)行此操作。按照說明操作，但使用以下設(shè)置：

Name: tensorrt-demo

GPU type: 環(huán)境 T4

Number of GPUs: 1

Machine type: n1-standard-2

如果你知道你使用的是大型模型，你可能需要一臺(tái)更強(qiáng)大的機(jī)器。

在 Boot disk 部分，選擇 CHANGE ，然后轉(zhuǎn)到 PUBLIC IMAGES 選項(xiàng)卡。對于 Operating system ，選擇 Linux 上的深度學(xué)習(xí) 。有很多版本，但請確保選擇 CUDA 版本。版本 基于 Debian 10 的深度學(xué)習(xí) VM 與 M98 適用于此示例。

其他設(shè)置可以保留為默認(rèn)值。

接下來， connect to the VM using SSH 。如果系統(tǒng)提示您安裝 NVIDIA 驅(qū)動(dòng)程序，請安裝。

在 VM 內(nèi)部，運(yùn)行以下命令以創(chuàng)建一些稍后使用的目錄：

mkdir models
mkdir tensorrt_engines

有關(guān)詳細(xì)信息，請參見 Create a VM with attached GPUs 。

建立形象

您需要一個(gè)自定義容器，其中包含執(zhí)行 TensorRT 代碼所需的依賴項(xiàng)： CUDA 、 cuDNN 和 TensorRT 。

您可以將以下示例 Dockerfile 復(fù)制到新文件中，并將其命名為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 。請記住，可能會(huì)有比本文所用版本更高版本的 Beam 可用。

通過在本地或在 GCE VM 中運(yùn)行以下命令來構(gòu)建映像：

docker build -f tensor_rt.dockerfile -t tensor_rt .

如果在本地執(zhí)行此操作，請執(zhí)行以下步驟。否則，您可以跳到下一節(jié)。

只有當(dāng)您在不同于您打算構(gòu)建 TensorRT 引擎的機(jī)器上創(chuàng)建圖像時(shí)，才需要以下命令。對于這篇文章，請使用 Google Container Registry 。將圖像標(biāo)記為用于項(xiàng)目的 URI ，然后推送到注冊表。確保用適當(dāng)?shù)闹堤鎿QGCP_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

創(chuàng)建 TensorRT 引擎

只有當(dāng)您在不同于要構(gòu)建 TensorRT 引擎的機(jī)器上創(chuàng)建圖像時(shí)，才需要以下命令。從注冊表中提取 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 中，您可以將其從本地計(jì)算機(jī)復(fù)制到/models 目錄：

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

現(xiàn)在，您應(yīng)該在 VM 中擁有 ONNX 模型和構(gòu)建的 Docker 映像?，F(xiàn)在是時(shí)候同時(shí)使用它們了。

以交互方式啟動(dòng) Docker 容器：

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

從 ONNX 文件創(chuàng)建 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

現(xiàn)在，您應(yīng)該可以在 VM 的/tensorrt_engines目錄中看到ssd_mobilenet_v2_320x320_coco17_tpu-8.trt文件。

將 TensorRT 引擎上傳至 GCS

將文件復(fù)制到 GCP 。如果您在將文件直接從 GCE 上傳到 GCS 時(shí)遇到gsutil 問題，您可能必須首先將其復(fù)制到本地計(jì)算機(jī)。

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

在 GCP 控制臺(tái)中，將 TensorRT 引擎文件上傳到您選擇的 GCS 存儲(chǔ)桶：

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

本地測試 TensorRT 引擎

確保您擁有使用 TensorRT RunInference 的 Beam 管道。一個(gè)示例是 tensorrt_object_detection.py ，您可以通過在 GCE VM 中運(yùn)行以下命令來遵循該示例。首先鍵入 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]"

您還創(chuàng)建了一個(gè)名為image_file_names.txt的文件，其中 包含圖像的路徑。圖像可以在 GCS 之類的對象存儲(chǔ)中，也可以在 GCE VM 中。

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

然后，運(yùn)行以下命令：

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

現(xiàn)在您應(yīng)該看到一個(gè)名為tensorrt_predictions.csv . 的文件。每行都有用分號分隔的數(shù)據(jù)。

第一項(xiàng)是文件名。

第二項(xiàng)是字典列表，其中每個(gè)字典對應(yīng)一個(gè)檢測。

檢測包含框坐標(biāo)（ ymin 、 xmin 、 ymax 、 xmax ）、分?jǐn)?shù)和類別。

有關(guān)如何在本地設(shè)置和運(yùn)行 TensorRT RunInference 的更多信息，請遵循 Object Detection 部分中的說明。

TensorRT Support Guide 概述了 GitHub 和產(chǎn)品包中所有支持的 NVIDIA TensorRT 8.5.1 示例。這些示例旨在展示如何在眾多用例中使用 TensorRT ，同時(shí)突出顯示界面的不同功能。這些示例在推薦器、機(jī)器理解、字符識別、圖像分類和對象檢測等用例中特別有用。

使用 DataFlow RunInference 運(yùn)行 TensorRT 引擎

現(xiàn)在您有了 TensorRT 引擎，就可以在 Dataflow 上運(yùn)行管道了。

下面的代碼示例是管道的一部分，您可以使用TensorRTEngineHandlerNumPy加載 TensorRT 引擎并設(shè)置其他推斷參數(shù)。然后讀取圖像，進(jìn)行預(yù)處理以將關(guān)鍵點(diǎn)附加到圖像，進(jìn)行預(yù)測，然后寫入 GCS 中的文件。

有關(guān)完整代碼示例的更多信息，請參見 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))

確保您已完成上一節(jié)中提到的 Google Cloud 設(shè)置。您還必須具有 Beam SDK installed 。

要在 Dataflow 上運(yùn)行此作業(yè)，請?jiān)诒镜剡\(yùn)行以下命令：

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"

根據(jù)模型的大小限制，您可能需要調(diào)整 machine _ type 、 GPU 的類型和計(jì)數(shù)或 disk _ size _ gb 。有關(guān)梁管道選項(xiàng)的詳細(xì)信息，請參見 Set Dataflow pipeline options 。

TensorRT 和 TensorFlow 目標(biāo)檢測基準(zhǔn)

為了進(jìn)行基準(zhǔn)測試，我們決定在前面提到的 SSD MobileNet v2 320 × 320 模型的 TensorRT 和 TensorFlow 對象檢測版本之間進(jìn)行比較。

在 TensorRT 和 TensorFlow 對象檢測版本中，每個(gè)推理調(diào)用都是定時(shí)的。我們計(jì)算了平均 5000 個(gè)推斷調(diào)用，由于延遲增加，沒有考慮前 10 個(gè)圖像。我們使用的 SSD 型號是小型型號。當(dāng)您的模型可以充分利用 GPU 時(shí)，您將觀察到更好的加速。

首先，我們將 TensorFlow 和 TensorRT 之間的直接性能加速與本地基準(zhǔn)進(jìn)行了比較。我們旨在證明 TensorRT 上降低精度模式的額外優(yōu)勢。

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

TensorRT FP32 的總體加速為 7.9x 。端到端包括數(shù)據(jù)副本，而 GPU 計(jì)算僅包括實(shí)際推斷時(shí)間。我們這樣做是因?yàn)槭纠Ｐ秃苄?。在這種情況下，端到端 TensorRT 延遲主要是數(shù)據(jù)拷貝。在更大的模型中使用不同的精度可以看到更顯著的端到端性能改進(jìn)，尤其是在推理計(jì)算是瓶頸而不是數(shù)據(jù)拷貝的情況下。

FP16 比 FP32 快 1.6 倍，沒有精度損失。 INT8 比 FP32 快 1.8 倍，但有時(shí)精度會(huì)降低，需要校準(zhǔn)過程。精度下降是特定于模型的，因此嘗試您的精度并查看產(chǎn)生的精度總是很好的。

使用 NVIDIA QAT 工具包的量化網(wǎng)絡(luò)也可以緩解此問題。有關(guān)詳細(xì)信息，請參見 Accelerating Quantized Networks with the NVIDIA QAT Toolkit for TensorFlow and NVIDIA TensorRT 和 NVIDIA TensorRT Developer Guide .

數(shù)據(jù)流基準(zhǔn)測試

在 Dataflow 中，使用早期實(shí)驗(yàn)中生成的 TensorRT 引擎，我們使用以下配置運(yùn)行：n1-standard-4 machine、disk_size_gb=75和 10 個(gè)工作人員。

為了模擬通過PubSub進(jìn)入 Dataflow 的數(shù)據(jù)流，我們將批大小設(shè)置為 1 。這是通過將ModelHandlers設(shè)置為最小和最大批量大小為 1 來實(shí)現(xiàn)的。

表 2 。數(shù)據(jù)流基準(zhǔn)

Dataflow runner 將管道分解為多個(gè)階段。通過查看包含推理調(diào)用的階段，而不是讀取和寫入數(shù)據(jù)的其他階段，可以更好地了解RunInference的性能。這在 Stage with RunInference 列中。

對于這個(gè)度量， TensorRT 只花費(fèi) TensorFlow 運(yùn)行時(shí)間的 57% 。如果你適應(yīng)一個(gè)完全使用 GPU 處理能力的更大模型，你預(yù)計(jì)加速度會(huì)增長。

度量推理_ batch _ latency _ micro _ secs 是對一批示例執(zhí)行推理所需的時(shí)間（以微秒為單位），即調(diào)用model_handler.run_inference的時(shí)間。這隨著時(shí)間的推移而變化，這取決于BatchElements的動(dòng)態(tài)批處理決策以及元素的特定值或dtype值。對于這個(gè)度量，您可以看到 TensorRT 比 TensorFlow 快 9.2 倍。

結(jié)論

在這篇文章中，我們演示了如何通過無縫拼接數(shù)據(jù)處理框架（ Apache Beam ）和推理引擎（ TensorRT ）來大規(guī)模運(yùn)行機(jī)器學(xué)習(xí)模型。我們提供了一個(gè)端到端的示例，說明如何將推理工作負(fù)載完全集成到數(shù)據(jù)處理管道中。

這種集成實(shí)現(xiàn)了一種新的推理流水線，該流水線有助于通過更好的 NVIDIA GPU 利用率和大大提高的推理延遲和吞吐量來降低生產(chǎn)推理成本。使用許多現(xiàn)成的 TensorRT 樣本，相同的方法可以應(yīng)用于許多其他推斷工作負(fù)載。未來，我們計(jì)劃進(jìn)一步自動(dòng)化 TensorRT 引擎構(gòu)建，并致力于 TensorRT 與 Apache Beam 的深度集成。