PyTorch→ONNXのコンバートでモデルの入力サイズを可変にして推論する

Posted On 2022-03-20

12.7k{icon} {views}

ONNXでモデルは入力サイズを可変にできます。PyTorch→ONNXの変換と、ONNXRuntimeでの推論方法、また可変にしたことによる速度の副作用を検証していきます。

きっかけ

モデルの推論高速化やエッジデバイスの展開において、ONNXへの変換は避けて通れないものですが、ONNXというと入力サイズ（shape）が固定のイメージがありました。ところが調べていると「可変の入力サイズでも変換できた」というのがあったので、それを検証していきます。

参考

まずは固定解像度から

まずはよくあるPyTorch→ONNXへの変換と、ONNXRuntimeでの推論を行ってみます。timmからEfficientNet-B0をダウンロードしてサクッとONNXへ変換してみます。

PyTorch以外にONNXとONNXRuntimeをインストールしておきます

pip install --upgrade onnx onnxruntime

import timm
import torch
import onnxruntime
import numpy as np

def convert_to_onnx_static():
    net = timm.create_model("efficientnet_b0")
    # 固定解像度のInput
    input = torch.randn(1, 3, 224, 224, requires_grad=True)
    # ダミーのテンソルを与えてexport
    # input_names, output_namesは入力層と出力層の名前（ONNXRuntimeで使うので設定しておく）
    torch.onnx.export(net, input, "efficientnet_b0_static.onnx", verbose=True, 
                      input_names=["input"],
                      output_names=["output"])

これを実行すると、次のようにずらーっと変換結果が表示されます。

中間層のShapeが表示されていますが、入力のShapeが(1, 3, 224, 224)の固定値なので全て実際の数値で記されています。これをONNXRuntimeで読みこんで推論するのは次のとおりです。

def read_onnx_static():
    session = onnxruntime.InferenceSession("efficientnet_b0_static.onnx")

    x = np.random.randn(1, 3, 224, 224).astype(np.float32)
    outs = session.run(None, {"input": x}) 

    print(outs[0], outs[0].shape) # outsはlist

    （中略）
   1.28479473e-06  5.53615905e-07 -4.19009302e-06 -3.03945376e-06
  -5.33000684e-06 -3.49969764e-06  9.33744377e-06  1.02329359e-05
 (1, 1000)

感覚的にはNumPyと同じ感覚で扱えます。ただONNX自体がプログラミング言語やフレームワークをまたいだ共通規格なので、他言語でも使いやすいというのが大きいですね。

このままでは可変サイズにできない

ここからがメインで、ONNXのモデルの入力サイズを可変にします。現状では、例えば入力サイズを(2, 3, 224, 180)などと変えるとエラーになります。

session = onnxruntime.InferenceSession("efficientnet_b0_static.onnx")

x = np.random.randn(2, 3, 224, 180).astype(np.float32)
outs = session.run(None, {"input": x})  # エラー
# onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Got invalid dimensions for input: input for the following indices
# index: 0 Got: 2 Expected: 1
# index: 3 Got: 180 Expected: 224
# Please fix either the inputs or the model.

これはPyTorch的にはOKですが、ONNXの作法的にはNGなのです。

入力サイズが可変のONNXモデルに変換する

PyTorch1.2から以下のようなdynamic_axesのオプションができました。関連issue。

例えば、先程のEfficientNet-B0を入力解像度とバッチサイズが可変のモデルにするには以下のようにします。

def convert_to_onnx_dynamic():
    net = timm.create_model("efficientnet_b0")
    # ダミーは固定のまま
    input = torch.randn(1, 3, 224, 224, requires_grad=True)
    # dynamic_axesのオプションを追加
    torch.onnx.export(net, input, "efficientnet_b0_dynamic.onnx", verbose=True, 
                      input_names=["input"],
                      output_names=["output"],
                      dynamic_axes={
                          "input": {0: "batch_size", 2: "height", 3:"width"}
                      })

「dynamic_axes」が新たに追加したオプションです。最初のキーの「input」は、この前に入力した引数の「input_name」の値に対応します。「input」の値は「インデックス：軸の名前」の順に指定します。名前はONNX向けのものなので何でも良いです。

ここではinputの0, 2, 3番目のインデックスが可変になります。変換のために入れる適当な値（乱数）は固定値でも、shapeは可変として取り扱われます。

可変のケースでの読み込み方

可変の場合でもそのまま読んでOKです。

def read_onnx_dynamic():
    # 可変なら変えてええやろで実はOK
    session = onnxruntime.InferenceSession("efficientnet_b0_dynamic.onnx")

    x = np.random.randn(10, 3, 180, 180).astype(np.float32) 
    outs = session.run(None, {"input": x}) 
    print(outs[0].shape) # (10, 1000)

意外と単純でした。ライブラリのバージョンによってはうまくいかないケースがあるかもしれません。

サイズを可変にして激重にならないの？

ONNXで何らかのコンパイルが入っていると、可変にしたとき遅いのではないか？　という疑問がわきます。

dynamic axesを指定したモデルで、固定 vs 可変

まずは、dynamic axesした可変のモデル（efficientnet_b0_dynamic.onnx）で、変換時の解像度で固定して推論したケースと、推論時の解像度をランダムに変えたケースを比較します。

各1000枚推論させ、10回試行したときの時間の最小値・中央値・最大値を見ます。コードは以下の通りです。

import time

def speed():
    session = onnxruntime.InferenceSession("efficientnet_b0_dynamic.onnx")

    result_static, result_dynamic = [], []
    for i in range(10):
        start_time = time.time()
        for j in range(1000):
            x = np.random.randn(1, 3, 224, 224).astype(np.float32)
            outs = session.run(None, {"input": x}) 
        result_static.append(time.time()-start_time)

    for i in range(10):
        start_time = time.time()
        for j in range(1000):
            x_res = np.random.randint(190, 240)
            y_res = np.random.randint(190, 240)
            x = np.random.randn(1, 3, x_res, y_res).astype(np.float32)
            outs = session.run(None, {"input": x}) 
        result_dynamic.append(time.time()-start_time)

    print(np.min(result_static), np.min(result_dynamic))
    print(np.median(result_static), np.median(result_dynamic))
    print(np.max(result_static), np.max(result_dynamic))

乱数の解像度は、CNNのおおよその計算量を目分量で考えてこの範囲にしました。正確な比較ではなく、要は可変で推論したときに激重にならないかを見るためのものです。結果は以下の通り。

10.851974487304688 10.755003213882446
11.24250078201294 11.08249819278717
11.385001420974731 11.287001132965088

ほぼ同じなので、dynamic_axesを指定したモデル同士なら、推論の解像度を変えても激重にはならないことがわかります。

固定のモデルと比較

次に推論解像度を固定し、Dynamic axesを指定しないモデル（efficientnet_b0_static.onnx）と比較を行います。

def speed_static():
    session = onnxruntime.InferenceSession("efficientnet_b0_static.onnx")

    result_static = []
    for i in range(10):
        start_time = time.time()
        for j in range(1000):
            x = np.random.randn(1, 3, 224, 224).astype(np.float32)
            outs = session.run(None, {"input": x}) 
        result_static.append(time.time()-start_time)

    print(np.min(result_static))
    print(np.median(result_static))
    print(np.max(result_static))

10.80599331855774
10.93899691104889
11.297000885009766

ほぼ誤差レベルでした。もしかしたら静的なサイズのほうが速いことはあるかもしれませんが、少なくとも可変にしたからといって激重になることはないことはわかります。

まとめ

PyTorch→ONNXにモデル変換し、ONNXRuntimeで可変の入力サイズを扱うことができた。また、推論速度も激重にはならなかった。

ONNXだけならpipから環境構築できて楽そうですし、個人的には結構ありな気がします。

環境

ONNX:1.11.0
ONNXRuntime：1.10.0
PyTorch：1.10.1+cu113、CPU推論

Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内

技術書コーナー

北海道の駅巡りコーナー

Tags:ONNX, ONNXRuntime, PyTorch, 推論, 高速化