RunInferenceでTensorRTを使用する

次の例は、BeamパイプラインでBERTベースのテキスト分類モデルを使用してRunInference APIとTensorRTを使用する方法を示しています。

推論のためのTensorRTエンジンの構築

Apache BeamでTensorRTを使用するには、トレーニング済みのモデルから変換されたTensorRTエンジンファイルが必要です。感情分析を行い、テキストを肯定的または否定的の2つのクラスに分類する、トレーニング済みの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エンジンができましたので、ローカルとGoogle Cloudの両方で実行できるBeamパイプラインでRunInferenceを使用してTensorRTエンジンを使用できます。

次のコード例は、パイプラインの一部です。`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ベンチマーク

ディスクサイズが75GBの`n1-standard-4`マシンを使用して、TensorRTエンジンと次の構成でDataflowで実験を実行しました。`PubSub`経由でDataflowにデータストリーミングを模倣するために、`ModelHandlers`の最小および最大バッチサイズを1に設定することで、バッチサイズを1に設定しました。

RunInferenceを含むステージ平均推論バッチレイテンシ（マイクロ秒）
T4 GPU搭載TensorFlow3分1秒15,176
T4 GPU搭載TensorRT45秒3,685

Dataflowランナーはパイプラインを複数のステージに分解します。RunInferenceのパフォーマンスをより正確に把握するには、データの読み書きを行う他のステージではなく、推論呼び出しを含むステージを確認します。これは「RunInferenceを含むステージ」列に示されています。

メトリクス`inference_batch_latency_micro_secs`は、バッチの例に対する推論の実行にかかる時間（マイクロ秒）です。つまり、`model_handler.run_inference`を呼び出す時間です。これは、BatchElementsの動的バッチ処理の決定と、要素の特定の値やdtype値によって時間とともに変化します。このメトリクスから、TensorRTはTensorFlowよりも約4.1倍高速であることがわかります。