torch2trt, torch_tensorrt の速度比較

どれが速いのか。

torch2trt
はPyTorchからTensorRTへのコンバーター。NVidiaがリリースしている。

torch_tensorrt
はtorchモデルをtensorrtにコンパイルできるライブラリ。torchがリリースしている。

両方tensorrtを使って高速化する。
どちらも簡単に使える。さて、速いのはどちらか。

結論

１、torch2trt、torch_tensorrtはtorchと比べて5倍以上速い。torchscriptよりも速い。
２、torch2trt、torch_tensorrtどちらが速いかはモデルによる。大差はない。
３、torch_tensorrtでしか変換できないモデルもある。VisionTransformerなど。

以上から、
僕はtorch_tensorrtを使います。

【EfficientNet】

	sec
torch	0.01614
torchscript	0.00689
tensorrt	0.00164
torch_tensorrt	0.00200

【Resnet】

	sec
torch	0.01101
torchscript	0.00646
tensorrt	0.00165
torch_tensorrt	0.00112

torch2trt

インストール

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

変換

import torch
from torch2trt import torch2trt
from torchvision import models
# create some regular pytorch model...
weights = models.ResNet50_Weights.IMAGENET1K_V1

model = models.resnet50(weights=weights).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])

推論

with torch.no_grad():
  y_trt = model_trt(x)

torch_tensorrt

インストール

pip install torch-tensorrt

変換

import torch_tensorrt
model_torch_trt = torch.compile(model, backend="tensorrt")
with torch.no_grad():
  y_torch_trt = model_torch_trt(x) # dry run

推論

with torch.no_grad():
  y_torch_trt = model_torch_trt(x)

🐣

---

フリーランスエンジニアです。
AIについて色々記事を書いていますのでよかったらプロフィールを見てみてください。

もし以下のようなご要望をお持ちでしたらお気軽にご相談ください。
AIサービスを開発したい、ビジネスにAIを組み込んで効率化したい、AIを使ったスマホアプリを開発したい、
ARを使ったアプリケーションを作りたい、スマホアプリを作りたいけどどこに相談したらいいかわからない…

いずれも中間コストを省いたリーズナブルな価格でお請けできます。

お仕事のご相談はこちらまで
[e-mail]

機械学習やAR技術を使ったアプリケーションを作っています。
機械学習／AR関連の情報を発信しています。

X
Medium
GitHub