ailia Tech BLOG
Keras から変換したONNXモデルに含まれる Transpose の除去
エッジ向け推論フレームワークであるailia SDKではONNXを使用してGPUを使用した高速な推論を行うことができます。この記事では、ailia SDKの開発の過程で得られたONNXのモデル最適化の知見を紹介します。
今回は ailia-models に掲載している hair_segmentation モデルの最適化手法を紹介します。
こちらのモデルは Keras で学習されたものを ax 株式会社にて ONNX 形式にコンバートしたものです。ONNX 形式モデルとして、Keras からコンバートしたそのままのファイルと、コンバート後に最適化を実行したものの二つが用意されています。
最適化内容
最適化前後の二つのモデルを Netron で表示して比較すると、次のようになります。
最適化前のモデルでは Clip の前後に Transpose が入っています。これらの Transpose のうち Clip の前に挿入されているのが channel first (N, C, H, W) のデータを channel last (N, H, W, C) に変換するもので、Clip の後に挿入されているものが channel last から channel first に戻すものです。
一方、最適化後のモデルでは Clip の前後にあった Transpose は削除され、Conv 直後の BatchNormalization も Conv と統合されています。
最適化の効果
ailia SDK 1.2.2 で backend に Vulkan を利用し、Windows 10 PC (GPU NVIDIA GTX 1650) 環境で 100 回推論を回した際の消費時間は、最適化前のモデルでは約 2.75 秒だったのに対して最適化後のモデルでは約 1.49 秒でした。約45%の処理時間削減に成功しています。
Kerasからの変換モデルで Transpose が多用される理由
ONNX の Conv や BatchNormalization は channel first (N, C, H, W) 形式を前提としているのに対して、Keras は channel last (N, H, W, C) 形式を採用するフレームワークです。
keras2onnx は ONNX モデルを出力する際、保守的に Conv / BatchNormalization の前後に Transpose を挿入して channel first / channel last の変換を行うことで channel order の相違を解消するため、Keras からの変換モデルには Transpose が頻出します。
最適化で Transpose を消せる理由
今回の最適化処理対象のコア部分を抽出すると次の画像になります。
そもそもこの二つの Transpose の間に Clip が挟まっていない場合であれば、変更して元に戻すだけの何もしない処理に該当するので、対消滅させてしまってもモデルの処理結果は変わらないことを理解いただけるはずです。
ここで考えているモデルは間に Clip が挟まっているとはいえ、この Clip は単一入力かつ単一出力で要素単位で処理するだけのノードです。channel first でデータを入力されても、channel last でデータを入力されても要素単位で考えれば結果には影響しません。つまり Clip を残して前後の Transpose を対消滅させてもモデルの処理結果は変わらないのです。
Transpose の削除を Python スクリプトで実装する方法
Transpose を対消滅させる場合、ノードの入力として供給されるパラメータを解釈する必要はなく、次の判定を行うだけで削除対象とすることが可能か判定できます。
- channel order 変更の Transpose であり、ノードの処理順の後ろで元に戻す Trnaspose も存在する
- 二つの Transpose の間に何も挟まっていないか、あるいは単一入力・単一出力の要素単位処理しか挟まっていない
この二つの条件を満たす Transpose をペアで削除対象として、後は他のノードの入出力を Transpose を削除しても矛盾がないように辻褄をあわせて繋ぎ合わせるだけです。前回の記事で解説した Conv と BatchNormalization の融合よりも容易に作成可能です。前回の内容を理解できた方であれば実装も簡単です。
まとめ
この記事では Keras から ONNX 形式に変換したモデルに含まれる Transpose の中には不要なものもあり、モデル内から不要な Transpose を削除することで推論の高速化が可能なことを紹介しました。
ailia SDKの無償評価版について
Transposeの除去を含むONNXの最適化スクリプトはailia SDKの無償評価版に同梱されています。ぜひ、ailia SDKの無償評価版をお試しください。
アイリア株式会社はAIを実用化する会社として、クロスプラットフォームでGPUを使用した高速な推論を行うことができるailia SDKを開発しています。アイリア株式会社ではコンサルティングからモデル作成、SDKの提供、AIを利用したアプリ・システム開発、サポートまで、 AIに関するトータルソリューションを提供していますのでお気軽にお問い合わせください。