PyTorchがONNXより早い場合について
こんにちは!Fusic 機械学習チームのトシュコフです。
機械学習モデルの開発から運用までなんでもしています。もし、機械学習で困っていることがあれば、気軽にお問い合わせください。
PyTorchが非常に使いやすく、モデルの実装がスムーズにできるライブラリーですが、モデルをデプロイする際の速度があんまり強みではありません。
関連コードを必要としないONNXというフォーマットがあります。
また、ONNXRuntimeを使用するとONNXモデルを早い速度で実行できます。
注意点がいくつものありますが、PyTorchモデルをONNX化した後、PyTorchモデルよりも遅い場合があります。
本記事ではPerfettoとONNXのプロファイリング方法を紹介します。
実験
Resizeとgrid_sampleのPyTorchとONNXの速度比較を行いました。以下のリンクでGoogle Colaboratoryで簡単に実行できます。
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でどのカーネルがどれくらいの時間がかかっているかを以下のように綺麗に表示できます。
まとめ
基本的にONNX RuntimeがPyTorchよりも早いですが、関数によって推論時間が遅くなることがあります。その場合、ボトルネックになっている所を調べるためにPrefettoを使う事をおすすめします。
最後に宣伝になりますが、機械学習でビジネスの成長を加速するために、Fusicの機械学習チームがお手伝いしています。機械学習のPoCから運用まで、すべての場面でサポートした実績があります。もし、困っている方がいましたら、ぜひFusicにご相談ください。
Discussion