torch2trtダウンロードtorch2trtソースコードのダウンロード

torch2trt

パイソン

v0.5.0

ダウンロード

TORCH2TRT

Tensortを使用して、どのモデルを使用していますか、または使用することを望んでいますか？ここでディスカッションにご参加ください。

TORCH2TRTは、Tensort Python APIを利用するPytorch to Tensort Converterです。コンバーターはです

使いやすい-1つの関数でモジュールを変換してtorch2trtを呼び出します
簡単に拡張しやすい - 独自のレイヤーコンバーターをPythonに書き込み、 @tensorrt_converterに登録します

問題が見つかった場合は、お知らせください！

このコンバーターは、Tensort / Pytorchのカバーが限られていることに注意してください。主にJetbotプロジェクトで使用されるモデルを簡単に最適化するために作成しました。コンバーターが他のモデルに役立つと思う場合は、お知らせください。

使用法

以下はいくつかの使用例です。詳細については、ノートブックをご覧ください。

変換する

 import torch
from torch2trt import torch2trt
from torchvision . models . alexnet import alexnet

# create some regular pytorch model...
model = alexnet ( pretrained = True ). eval (). cuda ()

# create example data
x = torch . ones (( 1 , 3 , 224 , 224 )). cuda ()

# convert to TensorRT feeding sample data as input
model_trt = torch2trt ( model , [ x ])

実行する

元のpytorchモデルと同じように、返されたTRTModuleを実行できます

 y = model ( x )
y_trt = model_trt ( x )

# check the output against PyTorch
print ( torch . max ( torch . abs ( y - y_trt )))

保存してロードします

モデルをstate_dictとして保存できます。

 torch . save ( model_trt . state_dict (), 'alexnet_trt.pth' )

保存されたモデルをTRTModuleにロードできます

 from torch2trt import TRTModule

model_trt = TRTModule ()

model_trt . load_state_dict ( torch . load ( 'alexnet_trt.pth' ))

モデル

test.shスクリプトを使用して、これらのモデルに対してコンバーターをテストしました。呼び出して結果を生成できます

./test.sh TEST_OUTPUT.md

以下の結果は、FPSのスループットを示しています。レイテンシを含む生の出力をベンチマークフォルダーに見つけることができます。

モデル	ナノ（pytorch）	ナノ（Tensort）	Xavier（Pytorch）	Xavier（Tensort）
alexnet	46.4	69.9	250	580
squeezenet1_0	44	137	130	890
squeezenet1_1	76.6	248	132	1390
resnet18	29.4	90.2	140	712
ResNet34	15.5	50.7	79.2	393
Resnet50	12.4	34.2	55.5	312
resnet101	7.18	19.9	28.5	170
ResNet152	4.96	14.1	18.9	121
densenet121	11.5	41.9	23.0	168
densenet169	8.25	33.2	16.3	118
densenet201	6.84	25.4	13.3	90.9
densenet161	4.71	15.6	17.2	82.4
VGG11	8.9	18.3	85.2	201
VGG13	6.53	14.7	71.9	166
VGG16	5.09	11.9	61.7	139
VGG19			54.1	121
vgg11_bn	8.74	18.4	81.8	201
vgg13_bn	6.31	14.8	68.0	166
vgg16_bn	4.96	12.0	58.5	140
vgg19_bn			51.4	121

設定

注：TORCH2TRTは、Tensorrt Python APIに依存します。 Jetsonでは、これは最新のJetpackに含まれています。デスクトップについては、Tensorrtインストールガイドに従ってください。また、デスクトップまたはジェットソン用のNGC Pytorch Dockerコンテナの1つにTorch2TRTをインストールしてみてください。

ステップ1 -TORCH2TRT Pythonライブラリをインストールします

TORCH2TRT Pythonライブラリをインストールするには、次のように電話してください

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install

ステップ2（オプション） - TORCH2TRTプラグインライブラリをインストールします

TORCH2TRTプラグインライブラリをインストールするには、次のように電話してください

cmake -B build . && cmake --build build --target install && ldconfig

これには、Tensortによってネイティブにサポートされない可能性のあるいくつかのレイヤーのサポートが含まれます。このライブラリがシステムに見つかると、TORCH2TRTの関連層コンバーターが暗黙的に有効になります。

注：TORCH2TRTは、CMAKEでコンパイルされた独立したライブラリとしてプラグインを維持するようになりました。これにより、コンパイルされたTensort Enginesがよりポータブルになります。必要に応じて、 python setup.py install --pluginsを呼び出すことにより、非推奨プラグイン（Pytorchに依存）は引き続きインストールされる場合があります。

ステップ3（オプション） - 実験的なコミュニティに貢献した機能をインストールします

TORCH2TRTを実験コミュニティで貢献したtorch2trt.contribをインストールするには、Quantization Aware Training（QAT）（ requires TensorRT>=7.0 ）など、以下に電話してください。

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt/scripts    
bash build_contrib.sh

これにより、ここにあるQAT例を実行できます。

どのように機能しますか？

このコンバーターは、変換関数（ convert_ReLUなど）を元のpytorch関数呼び出し（ torch.nn.ReLU.forwardなど）に接続することで機能します。サンプル入力データは、以前と同じようにネットワークに渡されます。登録された関数（ torch.nn.ReLU.forward ）が発生したときはいつでも、対応するコンバーター（ convert_ReLU ）がその後も呼ばれます。コンバーターは、元のPytorch関数の引数と返品ステートメント、および構築されているTensortネットワークに渡されます。元のPytorch関数への入力テンソルは、Pytorchテンソルに対応するTensortの対応物である属性_trtを持つように変更されます。コンバージョン関数は、この_trtを使用してTensortertネットワークにレイヤーを追加し、関連する出力テンソルの_trt属性を設定します。モデルが完全に実行されると、最終的なテンソルリターンはTensortネットワークの出力としてマークされ、最適化されたTensortエンジンが構築されます。

コンバーターを追加（またはオーバーライドする）方法

ここでは、Tensort Python APIを使用してReLUモジュールのコンバーターを追加する方法を示します。

 import tensorrt as trt
from torch2trt import tensorrt_converter

@ tensorrt_converter ( 'torch.nn.ReLU.forward' )
def convert_ReLU ( ctx ):
    input = ctx . method_args [ 1 ]
    output = ctx . method_return
    layer = ctx . network . add_activation ( input = input . _trt , type = trt . ActivationType . RELU )  
    output . _trt = layer . get_output ( 0 )

コンバーターは1つの引数、 ConversionContextを取得します。

ctx.network構築されているtensortネットワーク。
ctx.method_args指定されたpytorch関数に渡された位置引数。 _trt属性は、関連する入力テンソル用に設定されています。
ctx.method_kwargs指定されたpytorch関数に渡されたキーワード引数。
ctx.method_return指定されたpytorch関数によって返される値。コンバーターは、 _trt属性を関連する場合に設定する必要があります。

その他の例については、このフォルダーをご覧ください。