【Python高速化シリーズ】PyTorchとLibTorchを比較してみた②

こんにちは、AnyTechの赤川です。

前回の記事

tech.anytech.co.jp

この記事では前回に引き続き、PyTorchとPyTorchのC/C++版であるLibTorchとの間で比較検証を行い、LibTorchの採用がもたらすモデルの推論速度の改善度合いについて、考察しました。前回は一から定義したモデルについて速度比較をしましたが、今回はtorchvisionで提供される学習済みモデルを対象としました。学習済みモデルをTorchScript形式に変換し、それを利用にあたりPyTorchとLibTorchでの差異があるかを検証しました。

背景

例えばPyTorchで学習したモデルを他の言語で実装したいといった場合、以下のような選択肢があります。 - ONNXモデルに変換する - ONNXモデルを経由して別のフレームワーク（例えばtensorflowなど） - TorchScriptを利用する

ONNXモデルとは、モデルを表現するためのフォーマットです。ONNXを利用することで、PyTorchやtensorflowなど異なるフレームワークで開発されたモデルを単一のフォーマットで扱うことができるものになります。ONNXはランタイムも提供されているのでONNXに変換されたモデルをONNX上で実行することもできます。また、例えばPyTorchで作成したモデルをtensorflowで動かしたい場合は、PyTorchモデルをONNXモデルに変換し、そのモデルをさらにtensorflowモデルに変換するといったことも可能です。

ONNXモデルは便利で変換も簡単なのですが、今回は３つめのTorchScriptを対象とします。TorchScriptとはPyTorchで作成されたモデルを最適化し、Pythonに依存しない環境で呼び出したい場合などに利用されます（もちろん、TorchScriptに変換されたモデルをPyTorchで利用することもできます）。TorchScriptに関する詳細は本記事では省略させていただきますが、簡単な手続きを記述するだけで利用できます。詳しくは公式ページを参照ください。例として、torchvisionのResNet18の学習済みモデルをTorchScriptに変換する部分、およびそれを読み込んで使う方法についてサンプルコードを載せておきます。

PyTorchモデルのTorchScriptへの変換

import torch
from torchvision.models import resnet18

# ResNet18の学習済みモデルの取得
model = resnet18(pretrained=True)
# TorchScriptへ変換するためのダミーインプットの作成
dummy_input = torch.randn(1, 3, 224, 224)
# TorchScriptへ変換
traced_model = torch.jit.trace(model, dummy_input)
# モデルの保存
traced_model.save("resnet18.pt")

TorchScriptの読み込み（PyTorch)

import torch

# モデルの読み込み
model = torch.jit.load("resnet18.pt")

# 推論は通常のPyTorchモデルのように行えます
inputs = torch.randn(1, 3, 224, 224)
outputs = model(inputs)

TorchScriptの読み込み（LibTorch)

LibTorchを利用する場合、torch関連のヘッダーをインクルードする必要があります。PyTorchをインストールしている環境であれば、すでにPyTorchのインストール先に必要なヘッダーが保存されているので、そちらを参照するようにビルド時に指定してあげれば大丈夫です。cmakeを利用する場合は以下を参考に作成してください。

set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
list(APPEND CMAKE_PREFIX_PATH <site-packages内のtorchのルートフォルダパス>)
find_package(Torch REQUIRED)

add_executable(main main.cpp)
target_link_libraries(main "${TORCH_LIBRARIES}")
set_property(TARGET main PROPERTY CXX_STANDARD 17)

#include <torch/torch.h>
#include <torch/script.h>

int main(void){
    torch::jit::script::Module model = torch::jit::load("resnet18.pt");
    torch::Tensor input = torch::randn({1, 3, 224, 224});
    torch::Tensor output = model.forward({input}).toTensor();
}

検証条件

今回の検証では、torchvisionで利用可能な学習済みモデルをTorchScriptに変換した後、それらをPyTorchとLibTorchそれぞれで読み込み、推論速度について比較します。対象とするモデルは以下になります

実験環境
- MacBook Pro (13-inch, M1, 2020)
- Apple M1
- 8GBメモリ
検証モデル
- ResNet(18, 34, 50)
- VGG(11, 13, 16, 19)
- AlexNet
パフォーマンス測定方法
- バッチサイズを複数用意（1、8、16）
- 同じ条件で1000回実行し、その平均時間とする

検証結果

ResNet

ResNetについて検証を行った結果、以下の結果となりました。以下の結果を見ると、バッチサイズが小さい範囲では両者ほとんど差がないことが確認できます。ResNet18ではバッチサイズを大きくしても両者ほとんど差がないものの、構造を複雑にした場合にLibTorchのパフォーマンスの方が良いことが確認されました。特にResNet34の結果を見ると、バッチサイズを16にしたときにパフォーマンスがかなりさができていることが確認されました。リアルタイムにモデルを動かしたい（つまりバッチサイズが小さい）というユースケースであればどちらを使っても同じパフォーマンスが得られますが、まとめて一括処理したい場合はLibTorchに軍配が上がる結果となりました。

VGG

VGGについて検証を行った結果、以下の結果となりました。以下の結果を見ると、ResNetの時と同様にバッチサイズが小さい範囲では両者わずかな差しかないことが確認されます。VGG11とVGG13についてはバッチサイズを大きくしたときにわずかにLibTorchの方が処理速度が速いと言う結果になり、VGG16とVGG19では逆にPyTorchの処理速度の方が速いと言う結果になりました。特にVGG19について着目すると、バッチサイズの増加とともにLibTorchでのパフォーマンスが極端に悪化しています。測定するタイミングで別処理が走っているなどの理由で多少誤差はあるかと思いますが、ResNetの時とは対照的にPyTorchを使った方が良さそうな結果となりました。

AlexNet

AlexNetについて検証を行った結果、以下の結果となりました。この結果を見ると、バッチサイズが小さいときはPyTorchの方が、大きくするとLibTorchの方がパフォーマンスが良いと言う結果になりました。とはいっても、両者ほとんど差はないので、どちらを使ってもほとんど差がないと考えられます。

まとめ

今回は、学習済みモデルをTorchScriptに変換して、そのモデルをPyTorchとLibTorchで読み込むことでモデルのパフォーマンスに差が出るかを検証しました。結果としてはモデルに依存し、かつバッチサイズによって選択肢が変わりそうと言うことがわかりました。一方、マシンスペックなど今回検証条件に含めていない部分でパフォーマンスに影響を与える項目があるかと思うので、TorchScriptを使う際の参考にしてもらえればと思います。

AnyTech Engineer Blog

AnyTech Engineer Blogは、AnyTechのエンジニアたちによる調査や成果、Tipsなどを公開するブログです。