將 TensorRT 與 RunInference 搭配使用

NVIDIA TensorRT 是一個 SDK，可促進高效能機器學習推論。它旨在與 TensorFlow、PyTorch 和 MXNet 等深度學習框架搭配使用。它專注於最佳化和執行已訓練的類神經網路，以在 NVIDIA GPU 上有效率地執行推論。TensorRT 可以透過多種最佳化方式最大化推論輸送量，同時保留模型準確度，包括模型量化、層和張量融合、核心自動調整、多串流執行和有效率的張量記憶體使用。
在 Apache Beam 2.43.0 中，Beam 推出了 TensorRTEngineHandler，讓您可以在 Beam 管線中部署 TensorRT 引擎。RunInference 轉換簡化了 ML 推論管線建立流程，允許開發人員在生產管線中使用 Sklearn、PyTorch、TensorFlow，現在還可以使用 TensorRT 模型，而無需大量樣板程式碼。

以下範例示範如何在 Beam 管線中使用基於 BERT 的文字分類模型，將 TensorRT 與 RunInference API 搭配使用。

建立用於推論的 TensorRT 引擎

若要將 TensorRT 與 Apache Beam 搭配使用，您需要從已訓練模型轉換而來的 TensorRT 引擎檔案。我們採用基於已訓練 BERT 的文字分類模型，該模型會執行情緒分析並將任何文字分類為兩個類別：正面或負面。已訓練模型可從 HuggingFace 取得。若要將 PyTorch 模型轉換為 TensorRT 引擎，您必須先將模型轉換為 ONNX，然後再從 ONNX 轉換為 TensorRT。

轉換為 ONNX

您可以使用 HuggingFace transformers 程式庫將 PyTorch 模型轉換為 ONNX。如需詳細資訊，請參閱部落格文章使用 Hugging Face Optimum 將 Transformers 轉換為 ONNX。部落格文章說明需要安裝哪些套件。以下程式碼用於轉換。

from pathlib import Path
import transformers
from transformers.onnx import FeaturesManager
from transformers import AutoConfig, AutoTokenizer, AutoModelForMaskedLM, AutoModelForSequenceClassification


# load model and tokenizer
model_id = "textattack/bert-base-uncased-SST-2"
feature = "sequence-classification"
model = AutoModelForSequenceClassification.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# load config
model_kind, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model, feature=feature)
onnx_config = model_onnx_config(model.config)

# export
onnx_inputs, onnx_outputs = transformers.onnx.export(
        preprocessor=tokenizer,
        model=model,
        config=onnx_config,
        opset=12,
        output=Path("bert-sst2-model.onnx")
)

從 ONNX 到 TensorRT 引擎

若要將 ONNX 模型轉換為 TensorRT 引擎，請使用 CLI 中的下列命令

trtexec --onnx=<path to onnx model> --saveEngine=<path to save TensorRT engine> --useCudaGraph --verbose

若要使用 trtexec，請依照部落格文章使用 NVIDIA TensorRT 簡化和加速 Apache Beam 中的機器學習預測中的步驟。該文章說明如何從 DockerFile 建立可用於轉換的 Docker 映像檔。我們使用下列 Docker 檔案，該檔案與部落格文章中使用的檔案類似

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

FROM ${BUILD_IMAGE}

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

WORKDIR /workspace

RUN apt-get update -y && apt-get install -y python3-venv
RUN pip install --no-cache-dir apache-beam[gcp]==2.44.0
COPY --from=apache/beam_python3.8_sdk:2.44.0 /opt/apache/beam /opt/apache/beam

RUN pip install --upgrade pip \
    && pip install torch==1.13.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" ]

部落格文章還包含說明如何在本地測試 TensorRT 引擎的指示。

在 Beam 管線中使用 RunInference 執行 TensorRT 引擎

現在您有了 TensorRT 引擎，就可以在 Beam 管線中使用 TensorRT 引擎與 RunInference，該管線可以在本地和 Google Cloud 上執行。

以下程式碼範例是管線的一部分。您可以使用 TensorRTEngineHandlerNumPy 來載入 TensorRT 引擎並設定其他推論參數。

  model_handler = TensorRTEngineHandlerNumPy(
      min_batch_size=1,
      max_batch_size=1,
      engine_path=known_args.trt_model_path,
  )

  tokenizer = AutoTokenizer.from_pretrained(known_args.model_id)

  with beam.Pipeline(options=pipeline_options) as pipeline:
    _ = (
        pipeline
        | "ReadSentences" >> beam.io.ReadFromText(known_args.input)
        | "Preprocess" >> beam.ParDo(Preprocess(tokenizer=tokenizer))
        | "RunInference" >> RunInference(model_handler=model_handler)
        | "PostProcess" >> beam.ParDo(Postprocess(tokenizer=tokenizer)))

完整的程式碼可以在 GitHub 上找到。

若要在 Dataflow 上執行此作業，請在本機執行下列命令

python tensorrt_text_classification.py \
--input gs://{GCP_PROJECT}/sentences.txt \
--trt_model_path gs://{GCP_PROJECT}/sst2-text-classification.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-text-classification \
--sdk_container_image="us.gcr.io/{GCP_PROJECT}/{MY_DIR}/tensor_rt"

Dataflow 基準測試

我們使用 TensorRT 引擎在 Dataflow 中執行實驗，並採用下列組態：n1-standard-4 機器，磁碟大小為 75GB。為了模擬資料透過 PubSub 串流至 Dataflow，我們將批次大小設定為 1，方法是將 ModelHandlers 的最小和最大批次大小設定為 1。

	包含 RunInference 的階段	平均 inference_batch_latency_micro_secs
使用 T4 GPU 的 TensorFlow	3 分 1 秒	15,176
使用 T4 GPU 的 TensorRT	45 秒	3,685

Dataflow 執行器會將管線分解為多個階段。您可以查看包含推論呼叫的階段，而不是讀取和寫入資料的其他階段，來更好地了解 RunInference 的效能。這會在「包含 RunInference 的階段」欄中顯示。

指標 inference_batch_latency_micro_secs 是在範例批次上執行推論所需的時間 (以微秒為單位)，也就是呼叫 model_handler.run_inference 的時間。這會隨著時間變化，具體取決於 BatchElements 的動態批次決策，以及元素的特定值或 dtype 值。對於此指標，您可以看到 TensorRT 比 TensorFlow 快約 4.1 倍。

上次更新時間：2024/10/31

您找到您要的所有內容了嗎？

所有內容都有用且清楚嗎？您有任何想要變更的地方嗎？請告訴我們！