🧠

PyTorchがONNXより早い場合について

2023/07/07に公開

こんにちは!Fusic 機械学習チームのトシュコフです。

機械学習モデルの開発から運用までなんでもしています。もし、機械学習で困っていることがあれば、気軽にお問い合わせください。


PyTorchが非常に使いやすく、モデルの実装がスムーズにできるライブラリーですが、モデルをデプロイする際の速度があんまり強みではありません。
関連コードを必要としないONNXというフォーマットがあります。
また、ONNXRuntimeを使用するとONNXモデルを早い速度で実行できます。

注意点がいくつものありますが、PyTorchモデルをONNX化した後、PyTorchモデルよりも遅い場合があります。

本記事ではPerfettoとONNXのプロファイリング方法を紹介します。

実験

Resizegrid_sampleのPyTorchとONNXの速度比較を行いました。以下のリンクでGoogle Colaboratoryで簡単に実行できます。
Open In Colab

torchvisionのResizeとtorchのgrid_sampleが'nearest', 'biliear', 'bicubic'の3つモードがあります。

ランタイム 関数 Nearest Bilinear Bicubic
PyTorch Resize 0.86 ms 0.94 ms 0.91 ms
ONNX Resize 1.32 ms 1.32 ms 1.37 ms
PyTorch grid_sample 0.38 ms 0.34 ms 0.46 ms
ONNX grid_sample 5.11 ms 7.74 ms 15.19 ms

ONNXの速度調査方法

モデルで使用されている関数を1つずつONNX化することにより速度を計ることが困難のため、[ONNXRuntimeのProfiling Tools]ドキュメンテーションに書いている通り、sess_options.enable_profiling = Trueをする事により、Perfettoでどのカーネルがどれくらいの時間がかかっているかを以下のように綺麗に表示できます。

Prefetto Screenshot

まとめ

基本的にONNX RuntimeがPyTorchよりも早いですが、関数によって推論時間が遅くなることがあります。その場合、ボトルネックになっている所を調べるためにPrefettoを使う事をおすすめします。


最後に宣伝になりますが、機械学習でビジネスの成長を加速するために、Fusicの機械学習チームがお手伝いしています。機械学習のPoCから運用まで、すべての場面でサポートした実績があります。もし、困っている方がいましたら、ぜひFusicにご相談ください。

お問い合わせ

Fusic 技術ブログ

Discussion