🚝

WindowsでTensor-RT

2023/10/05に公開

Windows

ごあいさつ

QiitaはやっていましたがZennは初投稿です。
前々から書きたいと思っていたTensor-RTのインストール方法の解説を書きます。

1. Tensor-RTのインストールまで

一応解説しておきますと、Tensor-RTはニューラルネットワークの推論が2.5~5倍（僕調べ）になるめちゃくちゃつよーいライブラリです。
余裕があればこの辺のベンチマークなんかについても加筆していきたいと思いますが、ひとまずはインストール～PyTorchモデルを取り扱えるようにするまでをやります。

1-1. ダウンロード

Windowsの場合はNVIDIA公式サイトからダウンロードを行います。NVIDIAアカウントが必須なので所持していない方はあらかじめ作成しておきましょう。
バージョンは2023年10月時点で最新のTensor-RT 8.6 GA、Windows10向けのものを選択します。自分の使っているCUDAのバージョンにあったものを選ぶよう気を付けましょう。

1-2. 解凍

ZIPファイルの形式でダウンロードされるので展開します。ついでに、ダウンロードフォルダに置いたままにするのはなんだか気分が悪いので、私の場合はCドライブ直下に「DeepLearning」フォルダを作り、そこに移動させておきます（移動は必須ではありません）。

1-3. パスを通す

ダウンロードしたフォルダの中にある「lib」にパスを通します。
環境変数の編集画面で「Path」を選択。「DeepLearning」フォルダにTensor-RTを入れた私の場合は「新規」で「C:\DeepLearning\TensorRT-8.6.1.6\lib」を追加します。
または、libの中にあるDLLファイルをすべてCUDAの中（C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\vX.Y\bin）にコピってあげてもパスが通るらしいです。私は試していませんが、上記の方法でうまくいかなかった方は試してみるといいかもしれません。

1-4. インストール

まずは、Tensor-RTフォルダの中の「python」フォルダに移動します。続いて、自分の使用しているpythonバージョンに合ったwheelファイルをpipでインストールします（python.exe -m pip install tensorrt-*-cp3x-none-win_amd64.whl）

…と、文章で説明しても少々わかりにくいので具体例をあげて説明しましょう。
・ダウンロードしたTensor-RTのフォルダをCドライブ直下の「DeepLearning」フォルダに移動し
・Python3.9を現在利用しており
・TensorRT8.6.1.6のインストールを試みている

このような状況に置かれている私の場合は

cd "C:\DeepLearning\TensorRT-8.6.1.6\python"
python.exe -m pip install tensorrt-8.6.1-cp39-none-win_amd64.whl

これが打つべきコマンドになります。

続いて、付属する関連ライブラリをインストールする場合は引き続きpipコマンドを叩いていきます（ただ、どうやらこれはWindows&PyTorchの環境で使用する場合は不要らしいです。現に私はインストールせずとも問題なくTensorRT化を完遂できております）。

python3 -m pip install <installpath>\graphsurgeon\graphsurgeon-0.4.6-py2.py3-none-any.whl
python3 -m pip install <installpath>\uff\uff-0.6.9-py2.py3-none-any.whl
python3 -m pip install <installpath>\onnx_graphsurgeon\onnx_graphsurgeon-0.3.12-py2.py3-none-any.whl

エラーが出なければインストールは完了です。

2. PyTorchモデルを扱えるようにするまで

ひとまずTensor-RTのインストールが完了したわけですが、素のTensor-RTはonnxという~~めんどくさい~~規格のモデルしか対応していないので、どうにかして直接PyTorchのモデルを扱いたいところ。
ただし、onnxを介さずに学習済みモデルを扱えるライブラリをWindowsで使用しようとすると、大抵のものでは何らかの制限に引っ掛かります。

名称	Win対応
TF-TRT（後述）	×(?)
Torch-TensorRT	△
torch2trt	〇

たとえばTorch-TensorRTの場合はpipによるインストールができず、面倒なビルド作業を自力でこなす必要があります（詳しくはこちら、手順が少々複雑なので今回は深堀りせず、別の方法をとることにします）。
ちなみに、Torch-TensorRTとtorch2trtは何が違うんだ？という話ですが、~~私もわかりません~~どうやら前半はPyTorch主導でTensor-RTとPyTorchの連携を試みるプロジェクトに対し、後者はNVIDIA主導ということらしいです。
とどのつまり、私の調査が正しければWindowsユーザーの場合は必然的に選択肢は限られてくるようなので、torch2trtをこれからインストールしていきましょう。

2-1. ダウンロード

NVIDIA公式GitHubからダウンロードします。
gitを入れてない場合は右上の「code」と書かれた緑のボタンをおしてzipを落とせばOK。

2-2. 解凍

例のごとく解凍し、わかりやすい場所（私の場合は「DeepLearning」フォルダ）に移動します。

2-3. インストール

「torch2trt-master」フォルダに移動（私の場合「cd "C:\DeepLearning\torch2trt-master"」）し、インストールコマンド

python setup.py install

を打ち込むだけなのですが、ここで1つだけ大きな注意点があります。

それは、インストールしているCUDAバージョンとPyTorchの要求するCUDAバージョンが完璧に一致していなければならないということ。
PyTorchはTensorFlowと違い、CUDAバージョンの違いに鈍感なので、知らぬ間にバージョンがずれていたなんてこともあったりします。エラーの際には真っ先にここを点検すべきでしょう。
（筆者はこれでつまずきました。エラーメッセージ記録しておけばよかったかなぁ…）
エラーが出なければ終了です。お疲れさまでした。

3. torch2trtをスマートに使いこなす

TensorRT化の作業はかなりシンプルで、この10~20行程度の短いコードでこの工程は終了します。

import torch
from torch2trt import torch2trt
import network # 自分のネットワークを定義したファイル

MODEL_NAME="MyModel"
TRT_NAME="MyTRTModel"

BATCH_SIZE=1

with torch.no_grad():
    net=network.ResNet()
    net.load_state_dict(torch.load("model/"+MODEL_NAME+".pth"))
    net.cuda().eval().half() # fp16_modeの場合も、自分で.half()しておく必要はない（？）

    x=torch.rand((BATCH_SIZE,3,224,224)).cuda().half()

    model_trt=torch2trt(net,[x],fp16_mode=True,int8_mode=False,max_batch_size=BATCH_SIZE)

torch.save(model_trt.state_dict(),"trt_model/"+TRT_NAME+".pth")

しかしながら、少々厄介な仕様もいくつか存在するので、それを解説していきます。
まず第一に、何も手を打たないデフォルトの状態で、TensorRTは（というよりその裏側のPyTorchが？）勾配用のメモリを無意味に確保して盛大にメモリの無駄遣いをします。ただでさえメモリ消費の多いTensorRTでこれは致命的です。対策としてはwith torch.no_grad():で変換工程を囲って必ずや勾配をオフにしましょう。

次に、変換の工程で流したテンソルのサイズが推論解像度としてそのまま固定されます。また、この時のモデルの中間出力をもとにINT量子化の閾値などが決定されるようなので、INT8推論を行いたい方はこのような適当な乱数ではなく実データを流すべきでしょう。
推論時にどの精度を使用するかは引数で決定します。fp16とint8が両方Trueの場合はint8推論になる模様？

保存については通常のPyTorchモデルと同じ感覚で行えるようですが、ロードに関しては

from torch2trt import TRTModule

trt_net=TRTModule()
trt_net.load_state_dict(torch.load("trt_model/"+model_name+".pth"))

と、もうひと手間必要です。通常のPyTorchモデルのようにロードすることはできません。

また、これは私も最近知ったのですが、TensorRTは基本的にバージョン間での互換性がありません。つまり、もしもTensorRTのバージョンを入れ替えた場合は、TensorRT化を行うところからやり直しになります。
…はずだったのですが、TensorRT8.6以降からはこの制約が撤廃され、この面倒な再ビルド作業が不要になりました（ソース）
といっても、ここの記述を執筆中に正式リリースされている最新のTensorRTが8.6であるため、どのような仕様になるのかはいまいちよくわかっていないのですが…（次のアップデートではバージョンが8から9に上がるようなので、このようなメジャーアップデートでも本当に互換性が確保されるのかは謎です）。

まとめ

以上です。今回は、インストールからトラブルなくPyTorchモデルTensorRT化するまでの流れを書かせていただきました。
もうすぐfp8に対応したTensorRT v9やTensorRT-LLMの登場が予定されているようなので、これからももう少し情報を追記していくかもしれません。

4. おまけ・TensorFlowの場合

~~最近ではめっきり名前を聞かなくなってしまったTensorFlowですが~~、こちらでも一応TensorRTを利用できます。
具体的には、TensorFlow向けにカスタマイズされたTensorRTである「TF-TRT」を使うか、PyTorchに変換してしまうかの主に2つの選択肢があります。

4-1. TF-TRTを使う

実は、Windows版のTensorFlow2.8以降には、すでにこのTF-TRTが実験的に組み込まれています。組み込まれているというのは一見すると素晴らしいのですが、Windows版TF-TRTには不安要素も多く、手放しに持ち上げることはできません。
1つ目は、重ね重ねにはなりますが、実験的機能であること。公式ドキュメントやソースコード中でも実験的であることが強調されています。
2つ目は、TF-TRTはTensorFlow向けのカスタマイズの影響で素のTensorRTより遅い（と言われている）こと。本当の最速を目指したいならば、TF-TRTは使わない方がよいでしょう。
そして厄介なのが3つ目で、TensorFlowはv2.10.1を最後にWindowsのサポートが切られてしまいました（一応CPU対応のみintelが引き継ぐらしいですが、肝心要のGPUが使えないんじゃ…）。なので、間違っても「Windowsで手っ取り早くTensorRTしたいからPyTorchからTensorFlowに乗り換えよう！」みたいなのはおすすめできません。

どんなに頑張っても動きませんでした。TensorRTのバージョンを7.X（7.2.2以上8.4未満？）すれば動くみたいな未確認情報もありましたが、よくわからないので諦めました。

4-2. PyTorchへのモデル移植

前述の理由でTF-TRTを使用したくない場合は、PyTorchにモデルを移植してしまうという方法もあります。この記事の趣旨とは少々外れるので詳しい解説はしませんが、ここなんかを参考にすれば比較的簡単にTensorFlowモデルはPyTorchモデルに生まれ変わります。
移植さえ無事に済んでしまえば、あとはPyTorchのモデルをTensorRT化するときと全く同じ手順で作業は完了するでしょう。

参考

前半部分については、Tensor-RTのNVIDIA公式ドキュメントのうち、Windowsについて述べた3.2.4項（ここ）を大いに参考にさせていただきました。

環境

Windows11 Pro
CUDA 11.7
Python 3.9.16
Torch 2.0.1+cu117
Tensor-RT 未インストール → 8.6.1.6