ONNX RuntimeでGPUを指定する方法

ONNX RuntimeでGPUを指定する方法って知ってました ? もしかしたらほとんどの方が知っているのかも知れませんが、恥ずかしながら知りませんでした。というかできないと思ってました。

というのは、

ここ（How to choose CPU/GPU as the onnxruntime engine? #331）にある “For python, it’s not supported yet.” を見てしまったからでした。よく見てみれば、2019年初めの書き込みでしたね。onnxruntimeのバージョンも0.1.3とたいそう古いものになっています。

今ではバージョンは1.11.1まで上がって、もちろんのことながら、いろいろカスタマイズもできるようになってきました。本当に便利で重宝してます。

というわけで、

ONNX RuntimeでGPUを指定する方法ですが、とにかくまず公式のドキュメントを見てみましょう。

ONNX Runtime Execution Providers を見ると、さまざまなハードウェアがサポートされていることがわかります。もちろんAMDのグラボやARMもサポートされていますよ。機械学習はどうやってモデルを開発するかももちろん重要ですが、それをさまざまなデバイスにデプロイして運用できることもそれに負けず劣らず大切ですからね。ONNX Runtimeはそんな要求に簡単に対応できる貴重なプラットフォームの地位を確立しつつあると強く感じます。

ではさっそくサンプルコードを見てみましょう。

たとえばCUDAの場合だと、

CUDAのページに書いてある通り、こんな感じで書けます。

import onnxruntime as ort

model_path = '<path to model>'

providers = [
    ('CUDAExecutionProvider', {
        'device_id': 0,
        'arena_extend_strategy': 'kNextPowerOfTwo',
        'gpu_mem_limit': 2 * 1024 * 1024 * 1024,
        'cudnn_conv_algo_search': 'EXHAUSTIVE',
        'do_copy_in_default_stream': True,
    }),
    'CPUExecutionProvider',
]

session = ort.InferenceSession(model_path, providers=providers)

いや、素敵です。

前回の記事でOptionsの話しをしましたが、あれと組み合わせると、たとえばこんな感じになりますね。

options = onnxruntime.SessionOptions()
options.log_severity_level = 4
gpu = 1
session = onnxruntime.InferenceSession(onnx_path, options, providers=['TensorrtExecutionProvider', ('CUDAExecutionProvider', {'device_id': gpu}), 'CPUExecutionProvider'])

すばらすぃ～。

さいごに

公式のドキュメントを読めば済む話しでしたが、意外とさまよっている人がいるのではないかと思ったので、当たり前の記事を書いてみました。

ONNX Runtime本当にいろいろ助かってます。ひとつ四方山話しをすると、自作のLSTMがTensorflow Kerasで書いていたときはCPUでもかなり速度が速かったのですが、PyTorchに移行してみるとかなり遅くなってしまい、仕方がないのでGPUで動作させていました。これをONNXにConvertしてONNX Runtimeで動作させてみるとどうでしょう、またCPUでも速く動くようになったのです。

なぜ ? には深入りしない、場当たり的な話しで恐縮ですが、いろいろ助かっていて I love ONNX Runtime です。はい。

関連記事