common.title
Docs
Quantum Circuit
TYTAN CLOUD

QUANTUM GAMING

Desktop RAG
Overview
Terms of service

Privacy policy

Contact
Research
Sign in
Sign up
common.title

TensorRTを使ってみた

Yuichiro Minato

2023/11/14 14:44

こんにちは、最近ゆたぼんは学校に通ってるらしいです。

今日はTensorRTを見てみたいと思います。

Speeding Up Deep Learning Inference Using NVIDIA TensorRT (Updated)

https://developer.nvidia.com/blog/speeding-up-deep-learning-inference-using-tensorrt-updated/

NVIDIA TensorRTでディープラーニングの推論を高速化する

「NVIDIA TensorRTは、ディープラーニング推論用のSDKです。TensorRTは、すべての主要なディープラーニングフレームワークからトレーニング済みモデルをインポートするためのAPIとパーサーを提供します。その後、データセンターだけでなく自動車や組込み環境にも展開可能な最適化されたランタイムエンジンを生成します。」

推論用のSDKだそうです。

「TensorRTはC++とPythonの両方をサポートしています。いずれかを使用している場合、このワークフローの議論は役立つ可能性があります。」

C++とPythonをサポートしているみたいで、自分はPythonを使ってみます。

「CUDA対応GPUを搭載したコンピューター、またはGPUを備えたクラウドインスタンスとTensorRTのインストールが必要です。Linuxでは、NVIDIA Container Registry（NGC上）からTensorRT統合付きのGPUアクセラレーションされたPyTorchコンテナをダウンロードするのが最も簡単な方法です。」

ということで、ちょっと最新の利用はわからないので、早速TensorRT SDKのページを見てみます。

NVIDIA TensorRT

https://developer.nvidia.com/tensorrt

「NVIDIA® TensorRT™は、高性能ディープラーニング推論用のSDKで、ディープラーニング推論のための最適化ツールとランタイムを含み、推論アプリケーションにおいて低レイテンシと高スループットを実現します。」

とのことです。長所がいくつか並んでいるので、みてみます。

「推論を36倍高速化

NVIDIA TensorRTベースのアプリケーションは、推論時にCPUのみのプラットフォームより最大36倍高速に動作し、すべての主要なフレームワークで訓練されたニューラルネットワークモデルの最適化、高精度での低精度へのキャリブレーション、ハイパースケールデータセンター、組込みプラットフォーム、または自動車製品プラットフォームへのデプロイが可能になります。」

なんか心強いですね。

「推論性能の最適化

NVIDIA CUDA® 並列プログラミングモデルをベースに構築されたTensorRTは、量子化、レイヤーとテンソルの融合、カーネルチューニングなどの技術を使用して、NVIDIA GPU上での推論を最適化することを可能にします。」

量子化をはじめとしていろんなチューニングを行ってくれそうです。

「すべてのワークロードを加速

TensorRTは、量子化認識トレーニングおよびトレーニング後の量子化を使用したINT8および浮動小数点16（FP16）の最適化を提供し、ビデオストリーミング、レコメンデーション、不正検出、自然言語処理などのディープラーニング推論アプリケーションのデプロイを可能にします。精度を落とした推論は、多くのリアルタイムサービスや自律型および組込みアプリケーションに必要なレイテンシを大幅に低減します。」

学んだ量子化の技術の記述があります。効率化されそうです。

「Tritonを使用してデプロイ、実行、スケーリング

TensorRTで最適化されたモデルは、NVIDIA Triton™を使用してデプロイ、実行、スケーリングすることができます。Tritonは、TensorRTをバックエンドの一つとして含む、オープンソースの推論サービングソフトウェアです。Tritonを使用する利点には、動的バッチ処理や同時モデル実行による高スループットの実現、モデルアンサンブル、ストリーミングオーディオ/ビデオ入力などの機能が含まれます。」

なんかこちらはNVIDIAのAIプラットフォームの宣伝のようです。

「NVIDIA TensorRT-LLM

NVIDIA TensorRT-LLMは、NVIDIA GPU上で最新の大規模言語モデル（LLMs）の推論性能を加速し、最適化するオープンソースのライブラリです。開発者は新しいLLMを試すことができ、光速のパフォーマンスと迅速なカスタマイズをC++やCUDAの深い知識なしで実現します。

TensorRT-LLMは、TensorRTのディープラーニングコンパイラ（最適化されたカーネルを含むFasterTransformer、事前・事後処理、およびマルチGPU・マルチノード通信）をラップし、本番環境での推論のためにLLMを定義、最適化、実行するためのシンプルなオープンソースPython APIを提供します。」

流行りのLLMも対応しているようです。

早速ダウンロードしてみたいです。

NVIDIA TensorRT Download

https://developer.nvidia.com/tensorrt-download

完全に勘違いをしていました。本当にSDKなんですね。

NVIDIA TensorRTは、ディープラーニングアプリケーション向けの高性能ディープラーニング推論最適化ツールおよびランタイムです。

「TensorRTは、CUDAプラットフォームを使用するすべてのNVIDIA GPUで動作します。以下のファイルは、NVIDIA Quadro、GeForce、Tesla GPUを搭載したLinuxサーバーおよびワークステーション用です。NVIDIAは、本番環境へのデプロイメントにはTesla V100、P100、P4、P40 GPUの使用を推奨しています。」

SDKであり、APIであるということは、モデルや環境の最適化をしてテンソルの扱いをよしなにしてくれるということなんですね。なんかハードウェアに搭載された特殊な機能かと思っていましたが違うようです。ちょっと賢くなりました。

今回はチュートリアルにあった、TransformerモデルをTensorRTで構築します。

PythonAPIを活用します。インストールはpipでいけました。

!pip3 install -U torch torch-tensorrt tensorrt transformers tokenizers

完了したらまずはツールのimport

import torch
import torch_tensorrt
from transformers import BertModel

BERTモデルを読み込みます。インプットのテンソルサイズも決めておきます。

# Initialize model with float precision and sample inputs
model = BertModel.from_pretrained("bert-base-uncased").eval().to("cuda")
inputs = [
  torch.randint(0, 2, (50, 50), dtype=torch.int32).to("cuda"),
  torch.randint(0, 2, (50, 50), dtype=torch.int32).to("cuda"),
]

この辺りはオプションの設定値のようです。

# Enabled precision for TensorRT optimization
enabled_precisions = {torch.float}

Whether to print verbose logs

debug = True

Workspace size for TensorRT

workspace_size = 20 << 30

Maximum number of TRT Engines

(Lower value allows more graph segmentation)

min_block_size = 7

Operations to Run in Torch, regardless of converter support

torch_executed_ops = {}

TensorRTではモデルをコンパイルして最適化するみたいです。

# Define backend compilation keyword arguments
compilation_kwargs = {
  "enabled_precisions": enabled_precisions,
  "debug": debug,
  "workspace_size": workspace_size,
  "min_block_size": min_block_size,
  "torch_executed_ops": torch_executed_ops,
}

Build and compile the model with torch.compile, using Torch-TensorRT backend

optimized_model = torch.compile(
  model,
  backend="torch_tensorrt",
  options=compilation_kwargs,
)
optimized_model(*inputs)

optimized_modelができました。できるまでちょっと時間がかかります。

コンパイルした際と同じテンソルのサイズを作って新しい入力として入れてみます。

# Does not cause recompilation (same batch size as input)
new_inputs = [
  torch.randint(0, 2, (50, 50), dtype=torch.int32).to("cuda"),
  torch.randint(0, 2, (50, 50), dtype=torch.int32).to("cuda"),
]

そのまま作ったoptimized_modelに入れてみます。

import time
start = time.time()
new_outputs = optimized_model(*new_inputs)
print(time.time() - start)
print(new_outputs)

すぐ計算が終わりました。

0.041831254959106445
BaseModelOutputWithPoolingAndCrossAttentions(last_hidden_state=tensor([[[-0.3221, 0.5736, -0.0457, ..., -1.0612, 0.5637, -0.0993],
[-0.3019, 0.5389, -0.0299, ..., -1.0406, 0.5690, -0.0517],
[-0.2924, 0.5031, -0.0184, ..., -1.0449, 0.5474, -0.0179],
...,
[-0.2875, 0.5334, -0.0032, ..., -1.0319, 0.5706, -0.0367],
[-0.2879, 0.5241, -0.0096, ..., -1.0346, 0.5798, -0.0450],
[-0.2687, 0.5224, -0.0169, ..., -1.0363, 0.5742, -0.0369]],

    \[\[-0.4586, -0.2481, -0.3012,  ..., -0.4181,  0.7043, -0.1303\],
     \[-0.4141, -0.2181, -0.2709,  ..., -0.4549,  0.7023, -0.1575\],
     \[-0.4457, -0.2355, -0.2701,  ..., -0.4794,  0.6969, -0.1988\],
     ...,
     \[-0.4443, -0.2722, -0.2046,  ..., -0.4888,  0.7408, -0.1548\],
     \[-0.3940, -0.2922, -0.2079,  ..., -0.4660,  0.7510, -0.1474\],
     \[-0.3611, -0.2569, -0.2196,  ..., -0.4629,  0.7409, -0.1537\]\],

    \[\[-0.5010, -0.1904, -0.2949,  ..., -0.4588,  0.7204, -0.0918\],
     \[-0.4629, -0.1471, -0.2705,  ..., -0.4663,  0.7149, -0.1097\],
     \[-0.4521, -0.2143, -0.2668,  ..., -0.5133,  0.7339, -0.1611\],
     ...,
     \[-0.4613, -0.1858, -0.1967,  ..., -0.5278,  0.7775, -0.0908\],
     \[-0.4005, -0.1996, -0.1990,  ..., -0.5113,  0.7781, -0.0947\],
     \[-0.4080, -0.1479, -0.2186,  ..., -0.4997,  0.7590, -0.1075\]\],

    ...,

    \[\[-0.4645, -0.2301, -0.3143,  ..., -0.3653,  0.7413, -0.1582\],
     \[-0.4143, -0.1951, -0.2999,  ..., -0.4122,  0.7292, -0.1751\],
     \[-0.4015, -0.2322, -0.2817,  ..., -0.4374,  0.7278, -0.2052\],
     ...,
     \[-0.3906, -0.2465, -0.2104,  ..., -0.4444,  0.7840, -0.1915\],
     \[-0.3987, -0.2366, -0.2111,  ..., -0.4243,  0.7874, -0.1736\],
     \[-0.3911, -0.2276, -0.2301,  ..., -0.4152,  0.7794, -0.1814\]\],

    \[\[-0.4527, -0.1810, -0.2997,  ..., -0.4136,  0.7228, -0.1252\],
     \[-0.3693, -0.1850, -0.2896,  ..., -0.4672,  0.7320, -0.1405\],
     \[-0.3783, -0.1763, -0.2667,  ..., -0.4797,  0.7100, -0.1787\],
     ...,
     \[-0.3606, -0.1739, -0.2195,  ..., -0.4772,  0.7630, -0.1440\],
     \[-0.4160, -0.1853, -0.2271,  ..., -0.4693,  0.7631, -0.1235\],
     \[-0.3893, -0.1915, -0.2371,  ..., -0.4493,  0.7655, -0.1301\]\],

    \[\[-0.5075, -0.1582, -0.3554,  ..., -0.5009,  0.7158, -0.1251\],
     \[-0.4694, -0.1103, -0.3402,  ..., -0.4961,  0.7093, -0.1280\],
     \[-0.5144, -0.1393, -0.3394,  ..., -0.5319,  0.6910, -0.1664\],
     ...,
     \[-0.4557, -0.1573, -0.2680,  ..., -0.5801,  0.7467, -0.1091\],
     \[-0.4384, -0.1542, -0.2710,  ..., -0.5580,  0.7506, -0.1039\],
     \[-0.4258, -0.1298, -0.2880,  ..., -0.5397,  0.7450, -0.1225\]\]\],
   device='cuda:0', grad\_fn=<CompiledFunctionBackward>), pooler\_output=tensor(\[\[-0.3586, -0.4621, -0.3524,  ..., -0.5671, -0.6339,  0.3056\],
    \[ 0.0943, -0.3469, -0.8802,  ..., -0.9380, -0.5510, -0.2637\],
    \[ 0.0230, -0.3482, -0.8599,  ..., -0.9352, -0.5316, -0.1895\],
    ...,
    \[ 0.0511, -0.3516, -0.8612,  ..., -0.9357, -0.5331, -0.1930\],
    \[ 0.0351, -0.3530, -0.8705,  ..., -0.9360, -0.5479, -0.1842\],
    \[ 0.0116, -0.3540, -0.8269,  ..., -0.9195, -0.5529, -0.1753\]\],
   device='cuda:0', grad\_fn=<CompiledFunctionBackward>), hidden\_states=None, past\_key\_values=None, attentions=None, cross\_attentions=None)

もとのモデルも使ってみます。

start = time.time()
original_outputs = model(*new_inputs)
print(time.time() - start)
print(original_outputs)

なんかこっちの方が早いですが同じ答えが出ました。

0.01204538345336914
BaseModelOutputWithPoolingAndCrossAttentions(last_hidden_state=tensor([[[-0.3221, 0.5736, -0.0457, ..., -1.0612, 0.5637, -0.0993],
[-0.3019, 0.5389, -0.0299, ..., -1.0406, 0.5690, -0.0517],
[-0.2924, 0.5031, -0.0184, ..., -1.0449, 0.5474, -0.0179],
...,
[-0.2875, 0.5334, -0.0032, ..., -1.0319, 0.5706, -0.0367],
[-0.2879, 0.5241, -0.0096, ..., -1.0346, 0.5798, -0.0450],
[-0.2687, 0.5224, -0.0169, ..., -1.0363, 0.5742, -0.0369]],

    \[\[-0.4586, -0.2481, -0.3012,  ..., -0.4181,  0.7043, -0.1303\],
     \[-0.4141, -0.2181, -0.2709,  ..., -0.4549,  0.7023, -0.1575\],
     \[-0.4457, -0.2355, -0.2701,  ..., -0.4794,  0.6969, -0.1988\],
     ...,
     \[-0.4443, -0.2722, -0.2046,  ..., -0.4888,  0.7408, -0.1548\],
     \[-0.3940, -0.2922, -0.2079,  ..., -0.4660,  0.7510, -0.1474\],
     \[-0.3611, -0.2569, -0.2196,  ..., -0.4629,  0.7408, -0.1537\]\],

    \[\[-0.5010, -0.1904, -0.2949,  ..., -0.4588,  0.7204, -0.0918\],
     \[-0.4629, -0.1471, -0.2705,  ..., -0.4663,  0.7149, -0.1097\],
     \[-0.4521, -0.2143, -0.2668,  ..., -0.5133,  0.7339, -0.1611\],
     ...,
     \[-0.4613, -0.1858, -0.1967,  ..., -0.5278,  0.7775, -0.0908\],
     \[-0.4005, -0.1996, -0.1990,  ..., -0.5113,  0.7781, -0.0947\],
     \[-0.4080, -0.1479, -0.2186,  ..., -0.4997,  0.7590, -0.1075\]\],

    ...,

    \[\[-0.4645, -0.2301, -0.3143,  ..., -0.3653,  0.7413, -0.1582\],
     \[-0.4143, -0.1951, -0.2999,  ..., -0.4122,  0.7292, -0.1751\],
     \[-0.4015, -0.2322, -0.2817,  ..., -0.4374,  0.7278, -0.2052\],
     ...,
     \[-0.3906, -0.2465, -0.2104,  ..., -0.4444,  0.7840, -0.1915\],
     \[-0.3987, -0.2366, -0.2111,  ..., -0.4243,  0.7874, -0.1736\],
     \[-0.3911, -0.2276, -0.2301,  ..., -0.4152,  0.7794, -0.1814\]\],

    \[\[-0.4527, -0.1810, -0.2997,  ..., -0.4136,  0.7228, -0.1252\],
     \[-0.3693, -0.1850, -0.2896,  ..., -0.4672,  0.7320, -0.1405\],
     \[-0.3783, -0.1763, -0.2667,  ..., -0.4797,  0.7100, -0.1787\],
     ...,
     \[-0.3606, -0.1739, -0.2195,  ..., -0.4772,  0.7630, -0.1440\],
     \[-0.4160, -0.1853, -0.2271,  ..., -0.4693,  0.7631, -0.1235\],
     \[-0.3893, -0.1915, -0.2371,  ..., -0.4493,  0.7655, -0.1301\]\],

    \[\[-0.5075, -0.1582, -0.3554,  ..., -0.5009,  0.7158, -0.1251\],
     \[-0.4694, -0.1103, -0.3402,  ..., -0.4961,  0.7093, -0.1280\],
     \[-0.5144, -0.1393, -0.3394,  ..., -0.5319,  0.6910, -0.1664\],
     ...,
     \[-0.4557, -0.1573, -0.2680,  ..., -0.5801,  0.7467, -0.1091\],
     \[-0.4384, -0.1542, -0.2710,  ..., -0.5580,  0.7506, -0.1039\],
     \[-0.4258, -0.1298, -0.2880,  ..., -0.5397,  0.7450, -0.1225\]\]\],
   device='cuda:0', grad\_fn=<NativeLayerNormBackward0>), pooler\_output=tensor(\[\[-0.3586, -0.4621, -0.3524,  ..., -0.5671, -0.6339,  0.3056\],
    \[ 0.0943, -0.3469, -0.8802,  ..., -0.9380, -0.5510, -0.2637\],
    \[ 0.0230, -0.3482, -0.8599,  ..., -0.9352, -0.5316, -0.1895\],
    ...,
    \[ 0.0511, -0.3516, -0.8612,  ..., -0.9357, -0.5331, -0.1930\],
    \[ 0.0351, -0.3530, -0.8705,  ..., -0.9360, -0.5479, -0.1842\],
    \[ 0.0116, -0.3540, -0.8269,  ..., -0.9195, -0.5529, -0.1753\]\],
   device='cuda:0', grad\_fn=<TanhBackward0>), hidden\_states=None, past\_key\_values=None, attentions=None, cross\_attentions=None)

チュートリアルにあるように、テンソルのサイズを変更して入れてみたら、

# Does cause recompilation (new batch size)
new_inputs = [
  torch.randint(0, 2, (4, 14), dtype=torch.int32).to("cuda"),
  torch.randint(0, 2, (4, 14), dtype=torch.int32).to("cuda"),
]
new_outputs = optimized_model(*new_inputs)

再度コンパイルが始まりました。どうやら正しく動いているようには見えます。

もちょい大きなモデルで動かさないとあまり効果は見えないものなのかもしれませんが、ちょっと今回はわからないので、動いたのを確認できたので終わりにします。以上です。

© 2025, blueqat Inc. All rights reserved