サイバーエージェント AI Lab の Conversational Agent Team に所属している兵頭です。今回は私が半年ほど蓄積したONNXのチューニングテクニックを全てブログに残したいと思います。皆さんが既にご存知であろう基本的なことから、かなりトリッキーなチューニングまで幅広くご紹介したいと思います。長文になりますがご容赦願います。今回は応用編１です。

このブログのメインターゲット層は「リサーチャーが実装したモデルを実環境へデプロイするタスクを有する方々」です。一部リサーチャーの方々の参考になる情報が混じっていることもあるかもしれませんが、あまり興味を引かない内容だとは思います。リサーチャーメインの組織に属しながらリサーチエンジニアの立ち位置で身を投じていますので、研究の観点の少し手前あるいは少しその先の部分を担っている立場からこのブログを記載しているものとご認識願います。

いきなりですが、我々のチームでは一緒に Human Computer Interaction の研究・開発を行っていただけるリサーチエンジニアを募集しています。本ブログを見てご興味を持って頂けた方は是非一度カジュアルにお話させてください。よろしくお願いします。

• Cyberagent AI Lab

• Conversational Agent (対話エージェント 自動対話技術)

• 【AI事業本部・AI Lab】リサーチエンジニア（HCI)

  クリックするとYouTubeへ遷移します。

  

7. モデルのチューニング [応用]

基礎編 では、単一層のモデルやオペレーションの生成方法、簡易的なベンチマーク方法の一例をご紹介しました。今回は応用編として一歩踏み込んだモデルチューニングのテクニックをご紹介していきたいと思います。試行環境の構築方法やONNXの基本構造などをご理解いただいていることを前提として書き進めますので、基礎編 をご覧になっていない方はご一読いただいてからこちらの記事を読み進めていただくほうが理解が進みやすいと思います。この章では、simple-onnx-processing-tools に列挙されているツールをユースケースと利用方法を併せて上から順番にご紹介します。

7-1. モデルの融合 (snc4onnx)

２種類以上のモデルを１つのモデルへ融合します。これはあまり現実的なユースケースではないかもしれませんが、例えばVision系のモデル２種類であれば、１枚の２D画像を２つのモデルで１つの入力テンソルとしてシェアしつつ、１回の推論で２つのモデルの推論結果を同時に得られるようにすることができます。また、完成されたモデルではなくコンポーネントとしてバラバラにストックされたモデルを融合してひとつの end-to-end モデルへ構成することもできます。具体的に例を挙げてご紹介します。

分かりやすく２種類の論文実装モデルを１つのモデルに融合してみます。DMHead に動作イメージの動画やサンプルコード、実際に融合済みのONNXモデルをコミットしてあります。顔面あるいは頭部の６軸角度を推定するモデルです。イメージは下図のとおりです。

論文に記載されている定量的な評価は下表のとおりです。平均絶対誤差（MAE：Mean Absolute Error）つまり、MAE の数値が小さいほど実態との乖離が小さいということです。実際にワイルドな環境下で試さないと分からないことなのですが、SynergyNet は指標上はとても性能が高く見えますが、マスクを着用したり、実際にマスクを着用した画像でトレーニングを行うと著しく性能が劣化します。整った条件下でないと論文指標どおりの性能は一切出ません。ただ、ベースモデルが MobileNetV2 ですので超軽量です。一方で、6DRepNet はベースモデルの性能が高いためか、マスクを着用した顔画像を使用してトレーニングを行うと論文指標を上回る性能が出ます。VGG16 の改良版 RepVGG-B1g2 を使用しているため、重みの大きさとは裏腹に若干動作が速いですが、やはりトレードオフで MobileNetV2 に比べると動作が重たいと言わざるを得ません。また、6DRepNet と SynergyNet は正面を向いた顔しか推定できませんが、WHENet だけは後方を向いた後頭部でもある程度角度を推定できます。つまり、それぞれのモデルに一長一短があるため、それぞれのモデルのメリットのみを抽出したいというユースケースがあると仮定します。

1. 6DRepNet (RepVGG-B1g2) – 引用：https://github.com/thohemp/6DRepNet

   

   

2. SynergyNet (MobileNetV2) – 引用：https://github.com/choyingw/SynergyNet

   

   

3. WHENet (EfficientNet-b0) – 引用：https://github.com/Ascend-Research/HeadPoseEstimation-WHENet

   

   

下図のように、正面を向いた状態と後方を向いた状態の360°の範囲をできるだけ精度高く保ったまま６軸推定できるモデルを生成します。青色の枠は前向き、赤色の枠は後ろ向き、紫は中間ぐらい、というグラデーションで向いている方向が視認しやすいように表現しています。また、青の棒は Yaw、赤の棒は Roll、緑の棒は Pitch の角度を表します。これらのモデルは物体検出による顔位置の特定までは行わず、あくまで与えられた顔画像の６軸角度のみを推定するモデルですので、顔位置の検出には別途軽量な物体検出モデルを使用する必要があります。サンプル画像は YOLOv7-tiny を頭部画像データセットを使用してセルフトレーニングしたカスタムモデルを使用して頭部検出しています。全てのコンポーネントをNバッチ化済みですので、RAMの上限を超えない限り物体検出モデルからの出力内容を１度に何人分でも同時に投入できます。サンプル画像は１５人同時に検出していて、物体検出と顔向き推定と全ての前処理と後処理を合算して20ms弱です。TensorRTですので速いです。

では、さきほどご紹介したとおり、正面側の精度が高い 6DRepNet と後方を検出できる WHENet を snc4onnx を使用して融合します。

使用するモデル・コンポーネントの概要です。大きく分けて合計４つのコンポーネントを融合します。個々のオペレーションの生成方法は 基礎編 でご紹介済みですのでここでは取り上げません。画像がガタガタしてしまっている点はご容赦願います。

モデル・コンポーネント	概要
6DRepNet 157MB	N個の顔画像を受け取り、N個の正面側Yaw,Roll,Pitchを出力します。公式リポジトリの PyTorch のロジックを使用してONNXモデルを生成しました。
WHNet 17MB	N個の顔画像を受け取り、N個の全方位のYaw,Roll,Pitchを出力します。公式リポジトリの Keras のロジックを使用してONNXモデルを生成しました。生成難易度が高いですが、こちらに作業メモを残していますので気になる方はご覧ください。WHENetのエクスポート Pickleの形式に互換性が無くなってしまった Python3.5以前のレガシーな環境が必要です。なお、トレーニングコードが公開されていないため手軽に重みをチューニングすることはできませんが、ONNXに変換してしまえばこのブログのテクニックを使用してモデル構造を自由に書き換えることはできます。すでに非公式にモデル構造を改変してバッチ処理化した状態です。
shrunk_6drepnet_N	6DRepNetの出力のうちYawを-75°〜+75°の範囲のみ活性化するコンポーネントです。N個のYaw,Roll,Pitchを受け取り、YawをカットしたN個の出力を得ます。このコンポーネント自体も事前に snc4onnx を使用して複数のオペレーションをマージした状態です。
shrunk_whenet_N	WHENetの出力のうちYawを-76°〜-180°および+76°〜+180°の範囲のみ活性化するコンポーネントです。-75°〜+75°の範囲の角度を0°に補正します。N個のYaw,Roll,Pitchを受け取り、YawをカットしたN個の出力を得ます。このコンポーネント自体も事前に snc4onnx を使用して複数のオペレーションをマージした状態です。

２つのモデル、２つのコンポーネントをどのような順番で融合しても問題ありませんが注意点があります。どちらかというとコンポーネント生成時に注意すべき点です。

1. 融合したモデル内でオペレーションの名前が重複しないように全体で一意なネーミングにしておくこと
2. モデルのINPUT/OUTPUTの名前もオペレーションの名前と同様にモデル全体で一意なネーミングにしておくこと
3. 融合するAとBのコンポーネントの結合部は同じ形状なおかつ同じ演算精度にしておくこと

の３点です。
前置きが長くなりましたが融合してみます。ファイル名が長くパラメーターの切れ目が分かりくいですので、あえて改行とインデントを多めに入れて転記しています。

1. 6DRepNet と shrunk_6drepnet_N を融合

   snc4onnx \
   --input_onnx_file_paths \
       sixdrepnet_300w_lp_maskplus_prepost.onnx \
       shrunk_6drepnet_N.onnx \
   --srcop_destop \
       output_6drep_yaw_roll_pitch \
       shrunk_input \
   --output_onnx_file_path \
       sixdrepnet_300w_lp_maskplus_prepost_disable_rear_side_detection_Nx3x224x224.onnx

2. WHNet と shrunk_whenet_N を融合

   snc4onnx \
   --input_onnx_file_paths \
       whenet_Nx3x224x224_prepost.onnx \
       shrunk_whenet_N.onnx \
   --srcop_destop \
       yaw_roll_pitch \
       shrunk_input \
   --output_onnx_file_path \
       whenet_prepost_disable_front_side_detection_Nx3x224x224.onnx

3. 1.と2.を融合

   snc4onnx \
   --input_onnx_file_paths \
       whenet_prepost_disable_front_side_detection_Nx3x224x224.onnx \
       sixdrepnet_300w_lp_maskplus_prepost_disable_rear_side_detection_Nx3x224x224.onnx \
   --srcop_destop \
       input \
       input \
   --op_prefixes_after_merging \
       whenet \
       6drepnet \
   --output_onnx_file_path dmhead_merged_N.onnx

２種類のモデルを融合したため可視化すると横幅が大きくなりすぎて少々見づらいですが、左側が WHENet、右側が 6DRepNet のフローです。入力テンソルが２つのモデルで共用されている点と、出力の最終部で２つのモデルの出力を合算して推論結果を合成しています。単純加算で合成できる理由は少し前でご説明したとおり、２つのコンポーネントで出力の活性化と非活性化の角度範囲を制御していたからです。２つのモデルの出力角度がなるべく重複しないように（100%ではありません）調整されているため単純加算で乗り切れる状態になっていました。つまり、同じ角度範囲で２つのモデルからなるべく同時にゼロより大きい数値が出力されることが無いように調整したということです。最後の Add の部分は単純加算になっており境界付近では不安定になることが若干ありますので改善の余地があります。

• 結果

  6DRepNet はファインチューニングしてマスク顔に対応し、さらに論文性能より0.11ポイント精度を向上した自作モデルを使用しました。前後ともに検出力はまずまずのようです。

  6DRepNet: Yaw: 3.31, Pitch: 4.90, Roll: 3.36, MAE: 3.86
  WHENet:   Yaw: 5.11, Pitch: 6.24, Roll: 4.92, MAE: 5.42

  • 6DRepNet側で検出

    

  • シームレスにWHENet側で検出

    

snc4onnx のパラメータを転記するとブログが無用に長くなってしまいますので、応用編では詳細仕様が記載されたREADMEのURLのみご案内します。

snc4onnx の README はこちらです

ここまででモデルやコンポーネントの融合のテクニックをご紹介しましたが、このようなユースケースが活きると考えているのは、特定のタスクに特化した強力で軽量な部品単位のモデルを複数準備しておき、最終的には１つのモデルに融合して、実質マルチタスクに近いモデルをトレーニング後に自由に生成できるというところです。

7-2. モデルの分離 (sne4onnx)

１つのモデルを分離します。モデルの途中の特定のオペレーションの出力内容をデバッグするときや、7-1. モデルの融合 のようなユースケースでコンポーネントとして使用するためにモデルの一部分を抽出するとき、ハードウェアアクセラレータや推論フレームワークのバグを回避するためにモデルの一部分を別のコンポーネントに差し替える場合などがユースケースとして考えられます。また、世界中のリサーチャーが設計してくれたモデルは、実環境に投入して高効率に推論することまでを想定した綺麗なモデル構造になっていないことが多く、onnx-simplifier で最適化処理を施したときにAbortしてしまったり、モデルの最適化が十分に適用されないことが多いです。リサーチャーの観点からすれば、モデルを効率的に運用することや高パフォーマンスで推論することまでを想定しながらモデル設計することは本懐ではないと私自身は考えていますので、モデルの最適化タスクはエンジニアのタスクである、と割り切っています。リサーチャーが設計してくれた良いモデルを更に最適化するために、最適化が可能な範囲にモデルを分離して最適化を施し、のちに最適化済みのコンポーネント同士を再結合して全体を最適化したモデルを再構成する場合などが最もユースケースとして多いです。たとえば、一例でしかありませんしPyTorchのロジックをくまなく読み込んで自分で改修すれば問題ない事例ではありますが、Reshape の形状変更時に -1 が指定されたものが含まれているモデルはバッチ処理化をするときにAbortします。利用するエンジニア目線では -1 をできる限り指定してほしくないです。何故 Reshape に -1 が指定されている状態でバッチ処理化するとAbortするかという理由に触れておきます。Reshape後の形状指定が仮に [1,3,-1,512] だったとします。これをバッチ処理化するとどうなるかというと [-1,3,-1,512] となります。Reshape を行ったあとの出力の形状は分かりますでしょうか？ 無理です。つまり、ONNXに限らずどのフレームワークを使用したとしてもこの状況になるとAbortします。実質バッチ処理化して使えないモデルです。なお、-1 を指定した数値に書き換えるトリックは後続の節でご紹介します。

では、今回はとてもシンプルなモデルを分離してみます。MNIST – Handwritten Digit Recognition mnist-12.onnx

まず、onnx/models で公開されているサンプルONNXファイルをダウンロードし、基礎編でご紹介した onnxsim (onnx-simplifier) を３回実行してモデルの冗長さを排除します。

curl -OL https://github.com/onnx/models/raw/main/vision/classification/mnist/model/mnist-12.onnx
onnxsim mnist-12.onnx mnist-12.onnx
onnxsim mnist-12.onnx mnist-12.onnx
onnxsim mnist-12.onnx mnist-12.onnx

最適化後の MNIST ONNX の構造は下図のとおりです。

今回はモデルを分離してコンポーネント化する方法のご紹介のみにとどめるため、ユースケースに基づくモデルの分離ではない点にご注意願います。
では、下図の MaxPool の Pooling160_Output_0 の出力部分でモデルを分断してみます。分離をスタートするINPUT OPの名前と分離の終端となるOUTPUT OPの名前の両方を指定します。

sne4onnx \
--input_onnx_file_path mnist-12.onnx \
--input_op_names Input3 \
--output_op_names Pooling160_Output_0 \
--output_onnx_file_path mnist-12_part.onnx

分離されたONNXファイルが生成されました。Reshape と Gemm に含まれていた定数がモデルから除かれたかたちになりますのでファイルサイズが大きくシュリンクしているのが見て取れます。

Before	After

当然ながら、モデルの中間部分のみを分離することもできます。２つの MaxPool で挟まれた区間を分離してみます。

sne4onnx \
--input_onnx_file_path mnist-12.onnx \
--input_op_names ReLU32_Output_0 \
--output_op_names Pooling160_Output_0 \
--output_onnx_file_path mnist-12_part2.onnx

分離したONNXファイルが生成されました。 Conv に含まれていた定数がモデルから除かれたかたちになり、さらにモデルサイズがシュリンクしました。

Before	After

なおこのツールは、複数のINPUT NameとOUTPUT Nameを同時に指定することでひとつのモデルから複数の始点と複数の終点を持つモデルを抽出できるようにしてあります。 ONNXの標準API にも同様の処理が実装されていますし、onnx-graphsurgeon にも同様のインタフェースが用意されていますが、Pythonのソースコードを駆使してロジックを組み立てる必要があるため少々手間です。ちなみに、Idein さんも同じ思想で同様のツールを提供されています。onnigiri “The purpose of this package is to create subgraphs by partitioning computational graphs in order to facilitate the development of applications.” 私は完全にひとりでツールを作成していますので、Ideinさんのようにお仕事で作成されているツールをご利用いただくほうが運用の継続性の観点からはメリットがあると考えています。

sne4onnx の README はこちらです

7-3. Nodeの削除 (snd4onnx)

ONNXの標準APIや onnx-graphsurgeon を使用したPythonコードで表現するときにはかなり手間が多く複雑な操作です。そもそも、削除 という行為は複数のプリミティブな操作を組み合わせて実現する必要があり、コード量が若干多くなります。具体的には、

1. モデルを指定のNodeの入出力位置で上下に分断してコンポーネント化する
2. 削除対象のNodeをモデルから削除する
3. 分断した上下部分のコンポーネントの終端と始端が同じ形状および同じ算術精度であることをチェックする
4. 1.で分断した上下のコンポーネントを削除後の断面の入出力変数で結合しなおす
5. 全てのコンポーネントを再結合して生成しなおしたモデルの構造が破綻していないかチェックする
6. モデルをONNXファイルとして再出力する

という流れです。意味不明ですよね。でも、正常にモデルからNodeを削除するために必要な手番はこれだけあります。無理にこの節でご紹介するツールを使用しようとするとうまくNodeが削除できないパターンが多いですが、シンプルな構造のモデルであれば直感的に実施できます。実際に試してみます。snd4onnx というツールを使用します。

非常に限られたユースケースしかありませんが、不必要な cast を削除したり、フレームワークのバグを回避するための下準備として特定のNodeを削除する場合、などが該当します。
前節で使用した mnist-12.onnx を再び使用します。

意味の無い操作ですが、Conv の次にある Relu を削除してみます。

snd4onnx \
--remove_node_names ReLU114 \
--input_onnx_file_path mnist-12.onnx \
--output_onnx_file_path mnist-12_del.onnx

ONNXファイルが生成されました。

削除対象に指定した Relu Node名 ReLU114 がモデルから削除され、Conv と MaxPool が直結した状態のモデルが生成されました。

簡単にNodeが削除されてしまいましたが、現実的なシチュエーションでは複数の入力があるNodeや複数の出力があるNodeを削除したいことが多いと思いますので、このツール単体で手軽に実施できるパターンはごく限られています。削除はできても再接合部の形状不一致のワーニングやエラーが頻繁に出ます。あまり手を抜かず、他のツールを組み合わせて慎重にモデルをチューニングしたほうが確実です。

snd4onnx の README はこちらです

7-4. モデルサイズの圧縮 (scs4onnx)

これは 基礎編 でご紹介した onnx-simplifier を使用して最適化したONNXモデルを更にサイズ圧縮するためのテクニックを検証するために独自に考えた手法ですので、あまりユースケースとしては適切ではありません。したがいまして、ツールの使用方法や効果をこのブログで詳しくご紹介することはあえてしませんが、気になる方は scs4onnx のREADMEをご覧ください。

詳細には触れませんが、このツールのコンセプトだけを簡単にご説明します。

1. onnx-simplifier で各種オペレーションを最適化・定数化する過程でモデルが肥大化することがある
2. モデルが肥大化する原因は定数に置き換えられた画像のパッチ化処理などに必要となる大量のINT64精度のインデックス値など
3. INT64 は 64bit で値を保持するため、定数化されたテンソルの 要素数 x 64bit の容量がモデルに埋め込まれる
4. したがって、INT64 で埋め込まれた定数を INT32 などのより低精度な定数へ置換することでモデルのサイズを圧縮する
5. INT32であれば50%、INT16であれば25%、INT8であれば12.5% まで圧縮できる
6. ONNXファイル内部で複数回登場する定数を共用定数としてひとつに集約する
7. ONNXファイルから指定した定数のみを外部ファイル.npyへ切り出してONNXファイルを軽量化できる

このツールで効果が出るモデルは限られていますが、念の為結果の画像だけ貼り付けておきます。最も効果が高かったモデルでは、1.8GB あったONNXファイルが 2.1MB まで圧縮できました。ただし、ONNX本体から巨大な定数を外部ファイル(.npy)に抽出した結果です。量子化は行いませんので一切精度劣化はしません。

• 297.8MB -> 67.4MB (.onnx)

  

• 1.8GB -> 2.1MB (.onnx) + 884.7MB (.npy)

  

ツールの仕様でちょっとだけ面白い点に触れさせて頂くとすれば、モデル内部の定数部分を xxx.npy というNumpyのndarray形式のバイナリファイルにエクスポートする機能があることです。下図のクラゲのように線が伸びている部分が xxx.npy ファイルから読み出したテンソルを入力するための入力OPです。xxx.npy へエクスポートするタイミングで自動的にモデルへ入力OPを追加します。クラゲのようにひとつの入力OP部から複数の線が伸びているのは同じ定数を使用するOPが線の本数分だけ存在することを表します。したがって、xxx.npy へエクスポートしたONNXの定数は 再利用されるOPの数 x 定数のサイズ だけ圧縮されます。実は、繰り返し登場する同じ定数を集約するだけでモデルファイルの全体サイズを圧縮する効果はかなり大きいです。xxx.npy ファイルへ抽出しなくても scs4onnx を実行するとデフォルトの動作で定数を自動集約してファイルサイズを圧縮したONNXファイルを再出力する仕様にしています。定数を集約する、の意味がいまいちイメージできない方は 基礎編 の 5-1. ONNXモデルの構造 をご覧いただき、定数の名前を複数のNode間で共用するだけで簡単に実現できるトリックであることをご確認ください。プログラムで変数や定数を使い回しすることと概念に違いはありません。initializer 部で定義した定数の rawData を複数のNodeから参照することができる仕組みです。使い方として成り立つのかどうかが分かりませんが、モデルに投入する定数を推論ごとに差し替えることもできますね。意味不明ですが。もしかすると、定数部分を外部ファイルに抽出できているので、外部ファイルに抽出された定数ファイルのみを暗号化して秘匿してしまうとか、イロイロと遊べそうな要素がありそうです。復号化ロジックをどのように隠蔽するのかは考えていません。

具体的な仕様が気になる方はこちらの scs4onnx の README をご覧ください。

7-5. Nodeの生成 (sog4onnx)

基礎編 の 5-3. Node生成 でご紹介しましたので応用編では説明を省略します。

sog4onnx の README はこちらです

7-6. Nodeの属性値や定数の書き換え (sam4onnx)

モデルのチューニングを行ううえで最も使用頻度が高くなるであろうツール sam4onnx をご紹介します。このツールでできることを直感的に把握していただくために、先に画像で結果をご覧ください。

1. Transpose の perm 属性の書き換え例

   Before	After

2. Mul に埋め込まれた定数の書き換え

   Before	After

3. Reshape に埋め込まれた shape 定数の -1 を 1 に書き換え

   Before	After

4. NonMaxSuppression に埋め込まれた score_threshold 定数の -Infinity を 0.5 に書き換え

   Before	After

本質的には類似のコマンドを実行するだけなのですが、1. と 3. と 4. のコマンドを例示します。各Nodeの属性や変数・定数に指定可能な値の種類や精度はONNX公式の一覧表をご覧ください。 ONNX Operator Schemas

1. Transpose の perm 属性の書き換え例

   OP名 Transpose_17 の属性値 perm を [1,0] から [0,1] に書き換える

   sam4onnx \
   --input_onnx_file_path input.onnx \
   --output_onnx_file_path output.onnx \
   --op_name Transpose_17 \
   --attributes perm int64 [0,1]

2. Reshape の shape 定数の書き換え例

   OP名 Reshape_34 の入力定数 241 を [-1] から [1] に書き換える

   sam4onnx \
   --input_onnx_file_path input.onnx \
   --output_onnx_file_path output.onnx \
   --op_name Reshape_34 \
   --input_constants 241 int64 [1]

3. NonMaxSuppression の score_threshold 定数の書き換え例

   OP名 nonmaxsuppression11 の定数 score_threshold を [-Infinity] から [0.5] に書き換える

   sam4onnx \
   --input_onnx_file_path input.onnx \
   --output_onnx_file_path output.onnx \
   --op_name nonmaxsuppression11 \
   --input_constants score_threshold float32 [0.5]

このツールのメリットは、onnx-simplifier で最適化が終わった状態のONNXのパラメータを直接書き換えできることです。例えば、入力解像度に応じた複数種類のモデルを最適化した状態で用意しておき、後処理の検出閾値のみをあとから書き換えて検出精度の変化をテストしたい時や、前節で挙げたように、そもそも Reshape の変更後の形状に -1 が指定されている場合に onnx-simplifier の最適化やバッチ処理化が失敗する時に強制的に指定したパラメータに書き換えできることです。PyTorch や TensorFlow のロジックを一切読む必要がありません。また、おそらくモデル構造を最適化するタスクに取り組み始めるとすぐに直面する課題なのですが、PyTorch や TensorFlow からONNXを出力、あるいは変換して出力するときに問題になることがあるのが、PyTorch や TensorFlow のモデルオプティマイザが必要のないモデル構造変更を行ったり、少しの最適化を勝手にしてしまうことがあり、想定通りの構造のONNXモデルを素直に生成できないことがあります。そういう場合にも onnx-simplifier で最適化後のONNXファイルのパラメータを直接書き換えできるところに大きなメリットがあります。また、opset間で互換性の有無が有る属性値の追加や削除ができます。

sam4onnx の README はこちらです

7-7. opsetの書き換え (soc4onnx)

opsetの数字を強制的に書き換えるツールです。運用上はほとんど使いどころが無いため大幅に説明を省略します。下図のイメージです。opsetの数字を書き換えるだけでNodeの構造をopsetに対応した構造に書き換えることまではしません。ユースケースとしては、コンポーネント化したONNXファイルを複数結合する際にopsetアンマッチによる結合エラーを回避する場合が考えられます。

soc4onnx \
--input_onnx_file_path NonMaxSuppression.onnx \
--output_onnx_file_path NonMaxSuppression_13.onnx \
--opset 13

soc4onnx の README はこちらです

7-8. 入力OPのみのNCHW<->NHWC変換 および BGR<->RGB変換 (scc4onnx)

このツールもあまりユースケースが思い浮かびませんので説明を簡略化します。モデル全体を NCHW <-> NHWC 相互変換するのではなく、入力OPのみを NCHW <-> NHWC 相互変換します。また、BGR <-> RGB の変換用コンポーネントを入力OPの直後に外挿します。下図を見ていただくと雰囲気をご理解いただけると思います。

1. NCHW <-> NHWC 相互変換

   ONNXの内部はNCHW形式で処理されるため、入力OPをNHWC形式に変更した場合は、入力OPの直後で NCHW形式に転置し直す必要があります。あくまで推論用プログラムをONNXとTensorFlowで共用したいときにモデルの入力形式をNHWC形式へ統一して運用する、あるいはその逆のNCHW形式へ統一して運用する、ぐらいの使いどころぐらいしか思い当たりません。

   Before	After

2. RGB <-> BGR 相互変換

   入力OPの直後に３チャンネルをスライスして順番を入れ替える処理を外挿します。プログラムの前処理側でOpenCVを使用して定型的に実装することが多い、BGR -> RGB 変換 あるいはその逆の変換をモデルの入口に外挿してしまうイメージです。

   Before	After

scc4onnx の README はこちらです

7-9. Nodeの追加 (sna4onnx)

ONNXの標準APIや onnx-graphsurgeon を使用したPythonコードで表現するときにはかなり手間が多く複雑な操作です。7-3. Nodeの削除 と逆向きの操作を行うという点以外はほぼ同じです。追加 という行為は複数のプリミティブな操作を組み合わせて実現する必要があり、コード量が若干多くなります。具体的には、

1. モデルを指定のNodeの入出力位置で上下に分断してコンポーネント化する
2. 追加対象のNodeをsog4onnxを使用して生成する
3. 分断した上下部分のコンポーネントの終端と始端が2.で新規生成したNodeの入出力と同じ形状、同じ算術精度であることをチェックする
4. 1.で分断した上下のコンポーネントと2.で生成したNodeを間に挟んで入出力変数で結合しなおす
5. 全てのコンポーネントを再結合して生成しなおしたモデルの構造が破綻していないかチェックする
6. モデルをONNXファイルとして再出力する

こちらのツールも、無理に使用しようとするとうまくNodeが追加できないパターンが多いですが、シンプルな構造のモデルであれば直感的に実施できます。実際に試してみます。sna4onnx というツールを使用します。

前節で使用した mnist-12.onnx を再び使用します。

意味の無い操作ですが、Conv の次に Mul を追加してみます。

sna4onnx \
--input_onnx_file_path mnist-12.onnx \
--connection_src_op_output_name Convolution110 Convolution110_Output_0 dummy_mul inp1_var \
--connection_dest_op_input_name dummy_mul out1 ReLU114 Convolution110_Output_0 \
--add_op_type Mul \
--add_op_name dummy_mul \
--add_op_input_variables inp1_var float32 [1,16,14,14] \
--add_op_input_variables inp2_const float32 [1] \
--add_op_output_variables out1 float32 [1,16,14,14] \
--output_onnx_file_path mnist-12_add.onnx

ONNXファイルが生成されました。

追加指定した Mul Node名 dummy_mul がモデルに追加され、Conv と Relu の間を仲介しているモデルが生成されました。ただ、sna4onnx は定数を自動生成することができませんので、入力 B 側は inp2_const という名前の変数のみ定義された状態となっています。別途定数を定義して inp2_const に紐付ける必要があります。少し手間ですね。ややこしいですが、ONNXは Constant (定数) であってもひとつの Node (オペレーション) として独立して定義が存在しています。従って、Constant を新規に生成した場合は snc4onnx を使用して個別に融合する必要があります。ONNX公式の Constant オペレーション仕様はこちらに記載があります。Constant-12 一例としてopset=12の Constant を取り上げています。

定数を定義した Constant Node を生成します。

sog4onnx \
--op_type Constant \
--opset 12 \
--op_name const_custom1 \
--output_variables dummy_const float32 [1] \
--attributes value float32 [0.2]

生成されました。

では、前節で取り上げた snc4onnx を使用して Constant と Mul を結合してみます。

snc4onnx \
--input_onnx_file_paths Constant.onnx mnist-12_add.onnx \
--output_onnx_file_path mnist-12_merged.onnx \
--srcop_destop dummy_const inp2_const

融合済みのONNXファイルが生成されました。Conv と Relu の間に Mul を外挿することに成功しました。

少し手間ですが、任意の位置にNodeを追加することができました。オペレーションを１個追加する程度ではあまりメリットを感じませんが、大量に追加する場合はシェルで一気に操作できるため便利なときがあります。大半のシチュエーションではPyTorch側のロジックを調整して再エクスポートしたほうが早いです。

sna4onnx の README はこちらです

7-10. 固定バッチサイズを可変バッチサイズに書き換え (sbi4onnx)

可変バッチのONNXファイルを生成する場合、PyTorch側で torch.onnx.export(...) を実行するときに下記のように dynamic_axes にバッチサイズを表す次元を指定して初期化する文字列を与えるだけで生成することができます。下記の場合はゼロ次元目に batch_size と文字列が付与された入力OPが生成されます。

onnx_file = f"{MODEL}_HxW.onnx"
x = torch.randn(1, 3, 180, 320).cuda()
torch.onnx.export(
    self.model.module,
    args=(x),
    f=onnx_file,
    opset_version=11,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input' : {0: 'batch_size'},
        'output': {0: 'batch_size'},
    }
)
model_onnx = onnx.load(onnx_file)
model_onnx = onnx.shape_inference.infer_shapes(model_onnx)
onnx.save(model_onnx, onnx_file)

では、あえてこの章でご紹介する 固定バッチサイズを可変バッチサイズに書き換え のテクニックが必要となるシチュエーションはどんなものがあるかというと、

1. PyTorchモデルのソースコードが存在しない、あるいは提供されていない
2. バッチサイズが固定されたONNXは提供されているが可変バッチのONNXが提供されていない
3. リサーチャーが書いてくれたソースコードが複雑過ぎて読み込む作業の精神的ストレスが大きい

というパターンです。こういう状況はわりと多いです。従いまして、運用上は可変バッチのほうが扱いやすい場合を想定し、提供された固定バッチのONNXファイルを可変バッチのONNXファイルに書き換えます。

前節で使用した mnist-12.onnx を再び使用します。

mnist-12.onnx はゼロ次元目のバッチサイズが 1 に固定されています。ここを可変バッチの N に書き換えてみます。

sbi4onnx \
--input_onnx_file_path mnist-12.onnx \
--output_onnx_file_path mnist-12_Nbatch.onnx \
--initialization_character_string N

ONNXファイルが生成されました。ゼロ次元目のバッチサイズが指定した文字列 N で更新されています。

注意点として、このツールはゼロ次元目を数値から文字列へ強制的に書き換える動作のみを行います。したがいまして、ゼロ次元目を含む形状変更を伴う Reshape の形状変更、が含まれている場合や、バッチサイズそのものの意味を表現できないNodeが含まれている場合は書き換え動作が失敗します。シンプルに機械的にゼロ次元目を書き換える動作のみに対応しています。とはいえ書き換える方法が無いわけではなく、モデルのバッチサイズを書き換えたい場合は 7-2. モデルの分離 (sne4onnx) を使用して該当部分のNodeを切り離したうえで、切り離したコンポーネント個々の単位でバッチサイズを初期化し、バッチサイズ初期化後に再び 7-1. モデルの融合 (snc4onnx) を使用して再結合すれば問題ありません。既出のツールを組み合わせて作業手順を組み立てる発想の転換次第でほとんどの問題は回避可能です。このブログのテクニックを総動員して生成したモデルがこちらにあります。yolact_edge_onnx_tensorrt_myriad

sbi4onnx の README はこちらです

7-11. Node名の書き換え (sor4onnx)

入力OP、出力OP、中間Nodeの名称を書き換えることができます。ユースケースは、PyTorchから出力されたONNXのNode名が冗長なときに、運用ルールに則した分かりやすいネーミングに一括で統一する場合や、7-1. モデルの融合 (snc4onnx) を使用して複数のコンポーネントを結合する際、コンポーネント間でNode名の重複がある場合に、融合後のモデル全体で一意な名前になるようにあらかじめ名前を更新しておくような状況で良く利用します。ONNXのNode名はモデル全体で必ず一意になる必要があります。

では、再び mnist-12.onnx を使用します。

1. モデル全体のNodeを前方一致検索し Convolution を CONV に書き換え

   

   sor4onnx \
   --input_onnx_file_path mnist-12.onnx \
   --output_onnx_file_path mnist-12_rename1.onnx \
   --old_new "Convolution" "CONV" \
   --mode full \
   --search_mode prefix_match

   

2. 入力OPのみを前方一致検索し Input3 を input に書き換え

   

   sor4onnx \
   --input_onnx_file_path mnist-12.onnx \
   --output_onnx_file_path mnist-12_rename2.onnx \
   --old_new "Input3" "input" \
   --mode inputs \
   --search_mode prefix_match

   

3. 出力OPのみを前方一致検索し Plus214_Output_0 を output に書き換え

   

   sor4onnx \
   --input_onnx_file_path mnist-12.onnx \
   --output_onnx_file_path mnist-12_rename3.onnx \
   --old_new "Plus214_Output_0" "output" \
   --mode outputs \
   --search_mode prefix_match

   

名前の置換範囲や検索アルゴリズムには次のものが選択可能です。モデル全体、入力OPのみ、出力OPのみ、完全一致検索、部分一致検索、前方一致検索、後方一致検索。

-m {full,inputs,outputs}, \
--mode {full,inputs,outputs}

    Specifies the type of node to be replaced.
    full or inputs or outputs.
    full: Rename all nodes.
    inputs: Rename only the input node.
    outputs: Rename only the output node.
    Default: full

-sm {exact_match,partial_match,prefix_match,suffix_match}, \
--search_mode {exact_match,partial_match,prefix_match,suffix_match}

    OP name search mode.
    exact_match or partial_match or prefix_match or suffix_match.
    exact_match: Exact match search for OP name.
    partial_match: Partial match search for OP name.
    prefix_match: Prefix match search for OP name.
    suffix_match: Suffix match search for OP name.
    Default: exact_match

モデル構造を見直すときはかなり使用頻度が高くなるツールです。モデル利用者側から見た構造の視認性向上にも役立ちます。

sor4onnx の README はこちらです

7-12. モデルに出力OPを追加 (soa4onnx)

モデルの任意の位置のNodeから出力OPを追加生成します。すでにツール群 simple-onnx-processing-tools を活用している海外のエンジニアからのリクエストを受けて作成しました。 soa4onnx 各ノードの出力値をデバッグする場合や複数のモデルを融合するための接合点を生成するユースケースに使用します。

再び mnist-12.onnx を使用します。

最初の Relu から出力OPを生成してみます。加工前のモデルの出力は Plus214_Output_0 のひとつのみです。また、Relu の出力変数名は ReLU32_Output_0 です。

• モデルの全体構造

  

• 最初の Relu Nodeの情報

  

Relu の出力変数名 ReLU32_Output_0 をモデルへ追加するように指示します。

soa4onnx \
--input_onnx_file_path mnist-12.onnx \
--output_op_names "ReLU32_Output_0" \
--output_onnx_file_path mnist-12_output_added.onnx

ONNXファイルが生成されました。ReLU32 Nodeの出力 ReLU32_Output_0 がモデルの出力として追加され、モデル全体の出力が２つに増えました。

ちょっとしたモデル構造の改変程度ではあまり利用する機会がありませんが、複数のモデルを融合するときに接合点を生成するときに役立つときがあります。
soa4onnx の README はこちらです

7-13. モデルの出力OPを削除 (sod4onnx)

それほど多くはありませんが、PyTorchの公式論文実装から出力したONNXの出力OPに使用されないものが複数含まれていることがあります。無駄な出力を含めたままにしていると、ハードウェアアクセラレータとHostPC間での無駄なデータ転送コストが発生してレイテンシを悪化させることにつながることがあります。特にUSBを経由して使用するハードウェアアクセラレータなどを使用する場合はできる限りHostとやりとりする情報の転送量を抑えるのがレイテンシ低減に効果的です。なお、PyTorchの論文実装をカスタムして出力OPを削除する場合は不用です。

今回は複数出力OPを備えたモデルを使用してテストする必要がありますので、こちらの Robust Video Matting (RVM) モデルを使用します。

入力ひとつに対し、出力が fgr pha r1o r2o r3o r4o の６つあるモデルです。仮にこの出力のうち pha だけが必要とした場合のユースケースを想定します。

では、pha 以外の出力を全て削除してみます。

sod4onnx \
--input_onnx_file_path rvm_mobilenetv3_538x960.onnx \
--output_op_names "fgr" "r1o" "r2o" "r3o" "r4o" \
--output_onnx_file_path rvm_mobilenetv3_538x960_output_del.onnx

ONNXファイルが生成されました。

削除する、ということ自体は当たり前のようにできるのですが、注目していただきたい点は別の部分にあります。下図の出力OP削除前と出力OP削除後を交互にご覧ください。出力OP削除前に fgr に接続されていた各種OP (Add と Clip) が fgr 削除後に一掃され、最終出力に不必要なモデルのフローが全て自動的に削除されました。つまり、無駄な出力OPを全て削除することで通過しなくてもよい演算ルートを全て一掃する効果があるという点です。今回テストに使用したモデル rvm_mobilenetv3_538x960.onnx は削除された最終出力OPの手前に接続されているNodeが２つしかありませんでしたので推論速度向上に対する効果は限定的ですが、より一層複雑なモデルに対して同様の削除操作を行った場合、不必要な演算を大幅にカットできることがあります。PyTorchのロジックを目視で読み込んで手作業で最適化しなくてもツールが自動的に最適化してくれます。

• 出力OP削除前のモデルの最終出力近辺の様子

  

• 出力OP削除後のモデルの最終出力近辺の様子

  

sod4onnx の README はこちらです

7-14. モデルの中間Nodeの形状推定操作のみ実行 (ssi4onnx)

モデルのチューニングをするにあたって、モデルの中間Nodeの形状推定をすることのメリットは 基礎編 の 6-1-2. オペレーションの形状推定と定数変換と埋め込み でご紹介しました。しかしながら、状況によっては onnx-simplifier を使用してモデル構造の最適化はしたくないが各Nodeの入出力の形状のみ事前に推定しておきたいシチュエーションがあります。例えば、最適化までを行ってしまうとONNX以外の推論フレームワークで問題が発生する場合や、最適化前の形状を正しく把握してどのNode部分を最適化対象とするかを分析したいときなどです。onnx-smiplifier を使用すると自動的に形状推定を行いつつモデル構造を最適化してしまいます。

したがって、そういうユースケースに対応するため、このツール ssi4onnx はモデル構造の最適化を一切行わず、全Nodeの形状推定処理のみを実行します。

今回はこちらのモデル efficientnet-lite4-11.onnx を使用します。モデルの構造は下図のとおりです。onnx-simplifier で未最適化の状態です。

比較のためにまずは onnx-simplifier で形状推定とモデル構造の最適化の両方を実行してみます。

onnxsim efficientnet-lite4-11.onnx efficientnet-lite4-11_onnxsim.onnx

モデルの形状推定と構造の最適化後の様子です。４つ目の Conv の次にあった BatchNormalization が Conv に取り込まれて消滅するのと同時に各Nodeの入出力情報が書き込まれ、モデル構造の視認性が向上しました。

次に ssi4onnx を使用して形状推定のみ実行してみます。

ssi4onnx \
--input_onnx_file_path efficientnet-lite4-11.onnx \
--output_onnx_file_path efficientnet-lite4-11_inf_only.onnx

モデル構造の最適化が行われていないため BatchNormalization が Conv に取り込まれることなくそのまま残っていますが、各Nodeの入出力情報が書き込まれ、モデル構造の視認性が向上しました。

ssi4onnx の README はこちらです

7-15. ONNXの推論速度の簡易的なベンチマーク (sit4onnx)

基礎編 4. モデル生成の基礎 [PyTorch -> ONNX] で使用方法をご紹介しましたので応用編では説明を省略します。

7-16. ONNXのJSON変換 (onnx2json)

基礎編 5-1. ONNXモデルの構造 で使用方法や生成されたJSONファイルの読み方、IDEを使用した書き換え方法などをご紹介しましたので、応用編では説明を省略します。

onnx2json の README はこちらです

7-17. JSONのONNX変換 (json2onnx)

基礎編 5-1. ONNXモデルの構造 で onnx2json を使用して生成したJSONファイルに対して自由な編集を加えたモデル構造を再びONNXファイルへ逆変換するためのツールです。これまでの節でご紹介したツール群によるモデル加工の代わりに、IDEを使用してモデルの構造を直接手で書き換えてからONNXファイルへ戻すほうが早いシチュエーションが数多く存在します。たとえば、モデル内に繰り返し登場する類似の処理部分を一括で置換して書き換えたい場合などです。IDEの検索と置換の機能は高機能ですので、IDEの正規表現の置換機能などを駆使してテキストベースで一括変更すると楽な時があります。

ここでは 基礎編 5-1. ONNXモデルの構造 の説明をベースにしてモデルをIDEで書き換えたうえでONNXファイルへ戻すテクニックをご紹介します。

加工対象のサンプルとするモデルは 基礎編 5-1. ONNXモデルの構造 で生成した opset=17 の LayerNormalization2D_17.onnx です。事例はシンプルな方が分かりやすいと思いますので、定数のBase64エンコードが不要な属性値 epsilon を 1e-05 (0.000009999999747378752) から 1e-01 (0.10000000149011612) へ単純に書き換えます。

まずはONNXファイルをJSONに変換します。

onnx2json -if LayerNormalization2D_17.onnx -oj LayerNormalization2D_17.json

JSONファイルをVSCodeなどのIDEで開きます。編集箇所は "f": 1e-05, の部分です。編集したらJSONを上書き保存します。

• 編集前

  {
      "irVersion": "8",
      "producerName": "pytorch",
      "producerVersion": "1.14.0",
      "graph": {
      "node": [
          {
          "input": [
              "LayerNormalization2D_input",
              "/Constant_output_0",
              "/Constant_1_output_0"
          ],
          "output": [
              "LayerNormalization2D_output"
          ],
          "name": "/LayerNormalization",
          "opType": "LayerNormalization",
          "attribute": [
              {
              "name": "axis",
              "i": "-3",
              "type": "INT"
              },
              {
              "name": "epsilon",
              "f": 1e-05,
              "type": "FLOAT"
              }
          ]
          }
      ],
      "name": "torch_jit",
      "initializer": [
          {
          "dims": [
              "5",
              "10",
              "10"
          ],

• 編集後

  {
      "irVersion": "8",
      "producerName": "pytorch",
      "producerVersion": "1.14.0",
      "graph": {
      "node": [
          {
          "input": [
              "LayerNormalization2D_input",
              "/Constant_output_0",
              "/Constant_1_output_0"
          ],
          "output": [
              "LayerNormalization2D_output"
          ],
          "name": "/LayerNormalization",
          "opType": "LayerNormalization",
          "attribute": [
              {
              "name": "axis",
              "i": "-3",
              "type": "INT"
              },
              {
              "name": "epsilon",
              "f": 1e-01,
              "type": "FLOAT"
              }
          ]
          }
      ],
      "name": "torch_jit",
      "initializer": [
          {
          "dims": [
              "5",
              "10",
              "10"
          ],

次に、JSONファイルをONNXファイルへ逆変換します。

json2onnx -ij LayerNormalization2D_17.json -of LayerNormalization2D_17_eps.onnx

JSONファイルからONNXファイルが生成されました。ONNXファイルをNetronで表示すると、属性値 epsilon の値が 1e-01 (0.10000000149011612) へ更新されていることが確認できます。

このレベルの簡易的な修正であれば、前節までのツールを使用してスクリプトを実行するよりも、JSON化してテキストベースで編集を加えたうえでONNXファイルを再生成したほうが直感的で手っ取り早いことが多いです。使い慣れたIDEを使用してONNXを一括書き換えできます。

json2onnx の README はこちらです

7-18. 定数値のBase64エンコード／デコード (sed4onnx)

基礎編 5-1. ONNXモデルの構造 で扱い方をご紹介しましたので応用編では説明を省略します。

sed4onnx の README はこちらです

7-19. 簡易的な構造分析 (ssc4onnx)

基礎編 6-1-3. 使用するうえでのテクニック で扱い方をご紹介しましたので応用編では説明を省略します。

7-20. モデルの入出力形状名の書き換え (sio4onnx)

ONNXのモデルは未定義の次元 (unk__*) を受け入れます。例えば下記のような事例です。unk__* の部分はどのような形状でも受け入れることを表します。PyTorchからONNXをエクスポートする時に torch.onnx.export(...) で dynamic_axes に名前を指定しなかった場合などです。PyTorchの公式ドキュメントに記載があります。torch.onnx.export 下記の例では、yolonms_layer_1/* という名前の出力OPのゼロ次元目を未定義の次元に指定していますが、表示名を指定していませんので unk__* という少し長めで直感的ではない未定義の次元名が自動で付与されます。あるいは、onnx-simplifier を使用して形状推定とモデル構造の最適化を行っていない状態のときに形状未定という意味で unk__* が割り当てられることがあります。ONNXを推論で使用するうえでは特に問題は発生しませんが直感的なネーミングではなく、そして少し冗長です。このネーミング部分を sio4onnx で書き換えます。

dynamic_axes={
    "yolonms_layer_1/ExpandDims_1:0": [0,1],
    "yolonms_layer_1/ExpandDims_1:3": [0,1,2],
    "yolonms_layer_1/concat_2:0": [0,1],
}

ネーミングを書き換えます。unk__* の部分を batch や H や W や boxes などに書き換えるコマンドは下記のとおりです。面倒ですが、全ての入出力OP名を列挙する必要があります。前節のJSON変換を使用してIDEで書き換えても大して作業コストに変わりはありません。

sio4onnx \
--input_onnx_file_path yolov3-10.onnx \
--output_onnx_file_path yolov3-10_upd.onnx \
--input_names "input_1" \
--input_names "image_shape" \
--input_shapes "batch" 3 "H" "W" \
--input_shapes "batch" 2 \
--output_names "yolonms_layer_1/ExpandDims_1:0" \
--output_names "yolonms_layer_1/ExpandDims_3:0" \
--output_names "yolonms_layer_1/concat_2:0" \
--output_shapes 1 "boxes" 4 \
--output_shapes 1 "classes" "boxes" \
--output_shapes "boxes" 3

sio4onnx の README はこちらです

7-21. Node間コネクションの切り替え (svs4onnx)

7-9. Nodeの追加 (sna4onnx) および 7-3. Nodeの削除 (snd4onnx) では、任意のNode間に新たなNodeを外装したり、任意のNodeを削除する方法をご紹介しました。しかし同時に文中でも触れているように、外挿あるいは削除するNode位置の前後の入出力形状と演算精度が必ず一致していないとエラーになるということにも触れました。7-9. Nodeの追加 (sna4onnx) と 7-3. Nodeの削除 (snd4onnx) は内部処理が若干複雑な手番を踏んでいるがゆえに手軽に使える反面柔軟性に乏しい面がありました。そこでこのツール svs4onnx は少しだけ発想を変えて、コネクション間に新たなNodeを追加したり、コネクション間に存在するNodeを直接削除するのではなく、下記の順序で慎重に、そして正確にコネクションを接続替えする操作を数段階踏むことでNodeを外挿したり対象のNodeを削除する操作を表現するために作成しました。具体的な動作イメージは下記の通りです。

1. 入力ノードの出力変数に新規ノードを接合する
2. 新規ノードの出力変数をモデルの任意の入力変数へ接合する
3. 2.で新規Nodeの出力変数で入力変数が上書きされたNodeより手前の任意の数のNode(モデルの全体フローから外れたコンポーネント)をどのコネクションにも接続されていない離れ小島の状態にする
4. モデルの整合性をチェックし、全体構造から分離して離れ小島になっているコンポーネント部分をクリーニング(削除)する

わけのわからないことをゴチャゴチャと書きましたが、要は下図のとおりコネクションの接続替えをすることができる。ということです。この操作ができることで、7-9. Nodeの追加 (sna4onnx) および 7-3. Nodeの削除 (snd4onnx) の操作を表現できます。

Nodeの追加やNodeの削除というややこしい考え方は置いておいて、単純にこのツールでコネクションの接合先を変更するコマンドは下記の通りです。Cast Nodeの出力変数 cast1_output を Tile の入力変数 StatefulPartitionedCall/strided_slice_21 に接合するように指示を出しています。From から To へ接続せよ。というだけの指示です。このタイミングで、もともと接続されていたコネクションは自動的に消失します。

svs4onnx \
--input_onnx_file_path movenet_multipose_lightning_192x256_nopost_tmp1.onnx \
--from_output_variable_name "cast1_output" \
--to_input_variable_name "StatefulPartitionedCall/strided_slice_21" \
--output_onnx_file_path movenet_multipose_lightning_192x256_nopost_tmp2.onnx

svs4onnx の README はこちらです

7-22. ONNX(NCHW)形式からTensorFlow(NHWC)形式への自動コンバート (onnx2tf)

このツールだけは他のツールのようにモデルの構造を書き換える手段を提供するものではなく、シンプルにONNXをTensorFlowフォーマットへ変換するツールです。数年前からONNXの公式ツールとして onnx-tensorflow というものが提供されていたのですが、ツールを利用する上で大きな問題を抱えています。

1. NCHWをシミュレートするように無理やり変換するため、全ての Conv の前後に無駄な Transpose を大量に外装し、変換後のモデルの推論パフォーマンスが極端に悪化する
2. ロジックの構造が TensorFlow v1.x を前提としたとてもレガシーな作りになっている
3. 未定義の次元を含むモデルを変換できない
4. 内部の作りが複雑なことに加えて依存パッケージが多く正常に動作しないパターンが多い
5. コントリビュータが居ない

1.と2.の問題は大きな問題ではありますがOSSとして公開された当初からの仕様ですので、基本的な処理思想含めて大掛かりにリファクタリングをする必要があることを考えるとなんとか目をつぶることが出来ますが、5.はOSSとして致命的です。仮にプルリクエストを頑張って発行したとしてもレビューアが不在です。そして、内部のロジックの作りが複雑過ぎるためデバッグがとても困難でした。コントリビュータが36人も居て約5年？も継続的にメンテナンスされてきていたためとても残念ではありますが、メンテナ不在の時点で工数を掛ける気があまり湧きませんでした。

ということで、1.から5.の問題を全て解決すべく、ひとりでフルスクラッチしたツールがこちらのツール onnx2tf です。onnx-tensorflow が対応していない一部のオペレーションの変換にも追加で対応し、onnx-tensorflow で変換したモデルよりも 20% ほど高速なモデルを生成できます。ただし、RNN系とFFT系のオペレーションだけは TensorFlow v2.x 系での実装難易度が高すぎて今のところはうまく動作する状態で変換動作を実装することができていません。RNN LSTM GRU などです。もし気が向いた方がいらっしゃいましたら、是非プルリクエストをお送り下さい。

onnx2tf の README はこちらです

ちなみに、類似のツールとして onnx2tflite がありますが、現時点では以下の点で仕様の充足度が足りないと考えています。分かる範囲で記載しますので今後改善されていく可能性があるものも含みます。

1. tf.keras を多用しているため、saved_model を出力するときに頻繁にエラーが発生する

2. モデルの入力にひとつのテンソルしか許容しない (複数入力を許容しない)

3. 画像を入力にとるモデルしか考慮されていない

4. 量子化時のノーマライゼーションの実装に画像入力しか考慮されていない

5. per-channel の量子化しか考慮されておらず per-tensor の量子化ができないためハードウェアアクセラレータへの配慮が不足している

6. TPU や Myriad やその他のハードウェアアクセラレータに対するオペレーション置換のワークアラウンドが考慮されていない

7. ONNXからTFLiteへの変換しか考慮されておらず、その他のフレームワークへ転用することを考慮していない

8. 多次元の転置が考慮されていない

9. 汎用的な動作を定義する実装が無く、対応できるモデルの種類がとても少ない

10. TransposeConv の変換に問題がある

11. Slice の変換に問題がある

12. INT8量子化されたONNXモデルを変換することができない

13. 実装が不足しているオペレーションが多い (例えばNMSなど)

14. NHWC形式の入力を持つONNXモデルを変換できない

15. 変換エラーが発生するパターンを回避するための手段としては、プログラムを独自改造する方法しか無い

16. パッケージ化されていない

17. CLIインタフェースが無い

7-23. 名前が未定義のNodeに対する自動ネーミング (sng4onnx)

opset=8 より古いONNXのモデルにはNodeに名前が付与されていないことがあります。そういったモデルは onnx-simplifier やこのブログでご紹介したツール群でのモデル最適化に失敗することがあります。したがって、名前が付与されていないNodeに対して自動的に一意な名前を生成して書き込むツールを作成しました。そもそもONNXのモデルが古すぎることが問題であって、現時点ではあまり問題とはならないシチュエーションへの対応です。

名前を自動生成するONNXファイルと出力するファイル名を指定するだけですので細かい説明は省略します。

sng4onnx \
--input_onnx_file_path emotion-ferplus-8.onnx \
--output_onnx_file_path emotion-ferplus-8_renamed.onnx

結果は下図のとおりです。

sng4onnx の README はこちらです

7-24. doc_stringの自動クリーニング (sde4onnx)

こちらもかなり限られたシチェーションのみを想定したツールです。Hugging Faceで公開されている Stable Diffusion の vae_encoder.onnx をターゲットに作成しました。どう見てもゴミとしか思えない文字列の羅列を除去します。sergei さんの環境情報にはあまり興味がありませんので削除します。

sde4onnx \
--input_onnx_file_path vae_encoder.onnx \
--output_onnx_file_path vae_encoder_erased.onnx

sde4onnx の README はこちらです

7-25. GUIによるモデルのビジュアライズと加工 (OnnxGraphQt)

シェルやPythonのコードを書くことに抵抗を感じる方がいらっしゃるかもしれませんが、実はGUIで前述のツール群を組み合わせてモデル加工を行うためのOSSが誕生しています。 OnnxGraphQt というツールです。Netronを活用されている方々は多いと思いますが、このツールはビジュアライズをしたうえにモデル加工まで行うことができる点でアドバンテージがあります。バックエンドで私が作成したツール群をコールしていただいているため、私のツール群側のバグあるいはインタフェース不足にUI機能が制約されている可能性はあります。私のツール側、あるいは OnnxGraphQt へプルリクエストを発行してください。心よりお待ちしております。

前述でご紹介したようなシェルを実行する代わりにUIで直感的に操作することができるため馴染みやすいのではないかと思います。文字で説明するより画像で雰囲気をご覧ください。

• ONNXモデルの構造把握

• ディテールの確認

• Nodeの検索

トップ画面の右側ペインに、モデル加工を行うための操作ボタンが配置されており、私が作成したCUIツールとは違って、どのような操作ができるのか がボタンに明示されています。遷移先の子画面では、モデル加工に必要なパラメータをリストから選択したり直接キーボードで入力するようになっており、モデル加工を行ううえではとてもユーザーフレンドリーな設計になっています。ご興味がある方は是非お試し下さい。

X. おわりに

長くなりましたが、ONNXのモデルチューニングテクニック (応用編１) は以上です。次回は 応用編２ として、今回１記事で書ききれなかったテクニックをご紹介したいと思います。なお、次回の章タイトルは変更する可能性があります。

8. 各種PyTorchトリック

9. 各種フレームワーク向けモデルへの変換

   • TensorFlow
   • TensorFlow.js
   • TensorRT
   • CoreML
   • OpenVINO

サイバーエージェント AI Lab の Conversational Agent Teamに所属している兵頭です。今回は私が半年ほど蓄積したONNXのチューニングテクニックを全てブログに残したいと思います。皆さんが既にご存知であろう基本的なことから、かなりトリッキーなチューニングまで幅広くご紹介したいと思います。長文になりますがご容赦願います。

このブログのメインターゲット層は「リサーチャーが実装したモデルを実環境へデプロイするタスクを有する方々」です。一部リサーチャーの方々の参考になる情報が混じっていることもあるかもしれませんが、あまり興味を引かない内容だとは思います。リサーチャーメインの組織に属しながらリサーチエンジニアの立ち位置で身を投じていますので、研究の観点の少し手前あるいは少しその先の部分を担っている立場からこのブログを記載しているものとご認識願います。

いきなりですが、我々のチームでは一緒に Human Computer Interaction の研究・開発を行っていただけるリサーチエンジニアを募集しています。本ブログを見てご興味を持って頂けた方は是非一度カジュアルにお話させてください。よろしくお願いします。

• • Cyberagent AI Lab

  • Conversational Agent (対話エージェント 自動対話技術)

  • 【AI事業本部・AI Lab】リサーチエンジニア（HCI)

    クリックするとYouTubeへ遷移します。

    

1. このブログに記載すること

• ONNXモデル生成の基礎的な方法
• ONNXモデル生成の応用的な方法
• ONNXモデルの簡易的な推論方法
• ONNXモデルの簡易的なベンチマーク方法
• ONNXモデルのチューニング方法
• ONNXモデルのフレームワーク間コンバートのテクニック

2. このブログに記載しないこと

• ONNXとは何か
• モデルのトレーニング
• 数学的な演算ステップの効率化
• ランタイムに最適化したカスタムオペレーションの実装
• ONNX推論高速化のためのロジック側の工夫
• ONNX推論高速化のための onnxruntime / onnxruntime-gpu のパラメータチューニング
• ONNX以外のフレームワークの特徴
• ONNX以外のフレームワークを使用した推論方法
• ONNX以外のフレームワークへ転用したあとの推論の高速化手法

3. このブログに記載のソースコードを動かすための環境

このブログ内で使用する環境は下記のとおりです。HostPCの環境を壊さないために、 Docker や Google Colaboratory を活用するとトラブルなく試すことができます。 Anaconda は非推奨です。

• Ubuntu 20.04+ or Windows or Mac
• Python 3.8+
• PyTorch v1.13.0
• TensorFlow v2.10.0
• onnx v1.12.0
• onnxsim v0.4.8
• simple-onnx-processing-tools v1.0.76
• onnx2tf v1.1.25
• onnx-graphsurgeon 0.3.24

3-1. Docker を使用して試す場合

docker run --rm -it \
-v `pwd`:/workdir \
-w /workdir \
ghcr.io/pinto0309/onnx2tf:1.1.25

pip install torch torchvision torchaudio \
  --extra-index-url https://download.pytorch.org/whl/cpu

pip install onnxruntime==1.13.1

3-2. Google Colaboratory を使用して試す場合

!sudo add-apt-repository -y ppa:deadsnakes/ppa
!sudo apt-get -y update
!sudo apt-get -y install python3.9
!sudo apt-get -y install python3.9-dev
!sudo apt-get -y install python3-pip
!sudo apt-get -y install python3.9-distutils
!python3.9 -m pip install -U setuptools \
  && python3.9 -m pip install -U pip \
  && python3.9 -m pip install -U distlib
!sudo update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.7 1
!sudo update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 2
!python3.9 -m pip install tensorflow==2.10.0 \
  && python3.9 -m pip install -U onnx \
  && python3.9 -m pip install -U onnxruntime==1.13.1 \
  && python3.9 -m pip install -U nvidia-pyindex \
  && python3.9 -m pip install -U onnx-graphsurgeon \
  && python3.9 -m pip install -U onnxsim \
  && python3.9 -m pip install -U simple_onnx_processing_tools \
  && python3.9 -m pip install -U onnx2tf \
  && python3.9 -m pip install torch torchvision torchaudio \
    --extra-index-url https://download.pytorch.org/whl/cpu

3-3. HostPC を使用して試す場合 (PyTorchのバージョンは1.12.0以上であればどのバージョンでも問題ありません、GPUは不要です)

# CPU
pip install -U onnx \
&& pip install -U nvidia-pyindex \
&& pip install -U onnx-graphsurgeon \
&& pip install -U onnxsim \
&& pip install -U onnxruntime==1.13.1 \
&& pip install -U simple_onnx_processing_tools \
&& pip install -U onnx2tf \
&& pip install -U tensorflow==2.10.0 \
&& pip install -U torch torchvision torchaudio \
    --extra-index-url https://download.pytorch.org/whl/cpu

or

# CUDA 11.6 version
pip install -U onnx \
&& pip install -U nvidia-pyindex \
&& pip install -U onnx-graphsurgeon \
&& pip install -U onnxsim \
&& pip install -U onnxruntime-gpu==1.13.1 \
&& pip install -U simple_onnx_processing_tools \
&& pip install -U onnx2tf \
&& pip install -U tensorflow==2.10.0 \
&& pip install -U torch torchvision torchaudio \
    --extra-index-url https://download.pytorch.org/whl/cu116

4. モデル生成の基礎 [PyTorch -> ONNX]

良く知られたONNXを生成するための基本的なワークフローは、PyTorch で生成したモデルを torch.onnx.export(model, ...) あるいは torch.onnx.export(torch.jit.script(model), ...) で出力する方法です。改めてこのブログで説明するまでもない内容ではありますが、以降の章で説明する内容の前提知識として簡単に触れておきます。

ONNXを生成するために必要となる最低限の段取りは２つだけです。

1. PyTorch でモデルを定義する
2. PyTorch の torch.onnx.export(...) を実行してONNXファイルを生成する

PyTorch を使用して検証用の簡易的な１層モデルを生成します。複雑なモデルをサンプルとして生成するのは手間が掛かりますので、このブログでは LayerNormalization というオペレーションを１個だけ生成する PyTorch のソースコードを作成します。関数などでロジックを共通化したりせずプレーンに記載しています。モデルのトレーニングをすることを主眼としたブログではないため、本来は1.と2.の作業の間に モデルのトレーニング が必要であることにご注意願います。生成済みのモデルをチューニングする、という観点においては、モデルをトレーニングする、ということに関するテクニックはほとんど不要です。重みを微調整するのではなく、演算の順序や構成を組み替えて最適化することに主眼を置くためです。

from onnxsim import simplify は onnx-simplifier というONNXのモデル構造を機械的に最適化してくれるとても便利なツールをインポートしています。この章以降も頻繁に活用するため、あらかじめ作業用の端末にインストールしておく必要があります。

• make_LayerNormalization.py

  #! /usr/bin/env python
  
  import torch
  import torch.nn as nn
  import onnx
  from onnxsim import simplify
  import numpy as np
  np.random.seed(0)
  
  class LayerNormalization(nn.Module):
      def __init__(
          self,
          embedding_dim,
          weight,
          bias,
      ):
          super(LayerNormalization, self).__init__()
          self.embedding_dim = embedding_dim
          self.weight = weight
          self.bias = bias
  
      def forward(self, x):
          layernormed = nn.functional.layer_norm(
              input=x,
              normalized_shape=self.embedding_dim,
              weight=self.weight,
              bias=self.bias,
              eps=1e-05,
          )
          return layernormed
  
  if __name__ == "__main__":
      OPSET = [11, 17]
      MODEL = f'LayerNormalization2D'
      N, C, H, W = 20, 5, 10, 10
  
      for opset in OPSET:
          embedding_dim = [C, H, W]
          embedding_dim_tensor = torch.zeros(
              size=[C, H, W],
          )
          input = torch.randn(N, C, H, W)
          model = LayerNormalization(
              embedding_dim=embedding_dim,
              weight=torch.tensor(
                  data=torch.full_like(
                      input=embedding_dim_tensor,
                      fill_value=0.1,
                  ),
                  dtype=torch.float32,
              ),
              bias=torch.tensor(
                  data=torch.full_like(
                      input=embedding_dim_tensor,
                      fill_value=0.2,
                  ),
                  dtype=torch.float32,
              ),
          )
          onnx_file = f"{MODEL}_{opset}.onnx"
          torch.onnx.export(
              model,
              args=(input),
              f=onnx_file,
              opset_version=opset,
              input_names=[
                  f'{MODEL}_input',
              ],
              output_names=[
                  f'{MODEL}_output',
              ],
          )
          model_onnx1 = onnx.load(onnx_file)
          model_onnx1 = onnx.shape_inference.infer_shapes(model_onnx1)
          onnx.save(model_onnx1, onnx_file)
          model_onnx2 = onnx.load(onnx_file)
          model_simp, check = simplify(model_onnx2)
          onnx.save(model_simp, onnx_file)

作成したソースコードを実行して１層のみのONNXを生成します。

python make_LayerNormalization.py

LayerNormalization2D_11.onnx と LayerNormalization2D_17.onnx という２種類のONNXファイルが生成されました。これは、プログラム中の opset_version=opset の箇所でONNXのopsetバージョン番号に 11 と 17 を指定して２回実行しているためです。opset とは、ONNXのオペレーションのサブセットのようなものです。指定したバージョンごとに使用可能なオペレーションの種類が異なります。高いバージョンを指定すればするほど複数の処理をひとつのオペレーションに融合した効率的なオペレーションが割り当てられることがあります。詳しくはONNXの公式リポジトリにコミットされている こちら Operators.md の一覧をご覧ください。opsetごとに使用できるオペレーションに差が有ることが分かります。

では、opset=11 と opset=17 のそれぞれのバージョンで生成したONNXファイルにどのような差が生じたかをご覧ください。どちらも LayerNormalization の処理であり、入力の値が同じであれば出力の値は同じになります。従って、左の opset=11 で生成したプリミティブな演算の組み合わせを順次実行する処理が右の opset=17 で生成したたったひとつのオペレーションに融合していることが分かります。ONNXのモデルを onnxruntime という実行フレームワークへ読み込ませて実行したときには opset=11 であっても opset=17 であっても基本的には同じ振る舞いをします。この章は基本の章ですので詳細は割愛して後続の章で説明しますが、高いバージョン番号のopsetを指定してモデルの構造を綺麗にまとめたほうがメリットが大きいように見えるかもしれませんが、ONNX以外のフレームワークやTPUやMyriadなどの推論専用ハードウェアアクセラレータへモデルを転用するときに問題が生じることが多いです。そのため、あえて少しだけ古いopsetを指定してプリミティブな演算で構成されたオペレーションに変換しておく、というテクニックが活きるときがあります。実はこのテクニックが活きるシチュエーションはとても多いです。というのも、ハードウェアアクセラレータ側のランタイムのバージョンアップがONNX側の仕様改善のスピードに全く追いつけないからです。最たる例はTensorRTです。

• Netron でモデル構造を可視化

  Opset=11	Opset=17

では、生成された LayerNormalization2D_11.onnx と LayerNormalization2D_17.onnx が正しく動作することを確認してみます。しかしながら、検証用の推論コードは 書きません。ONNXの動作を簡易的に検証するためのツールを自作していますのでそちらを利用して動作を確認します。３章の手順を正しく行って環境構築が終わっている場合はすでにそのツールが導入済みの状態になっています。ONNXの簡易動作検証用のツールは sit4onnx と名付けてパブリックに公開しています。

使い方は下記のとおりです。ヘルプメッセージは全て英語で記述していますので、READMEの日本語チュートリアル化のプルリクエストを歓迎します。

• コマンドラインインタフェース

  sit4onnx -h
  
  usage: sit4onnx
  [-h]
  -if INPUT_ONNX_FILE_PATH
  [-b BATCH_SIZE]
  [-fs FIXED_SHAPES [FIXED_SHAPES ...]]
  [-tlc TEST_LOOP_COUNT]
  [-oep {tensorrt,cuda,openvino_cpu,openvino_gpu,cpu}]
  [-ifp INPUT_NUMPY_FILE_PATHS_FOR_TESTING]
  [-ofp]
  [-n]
  
  optional arguments:
    -h, --help
      show this help message and exit.
  
    -if INPUT_ONNX_FILE_PATH, --input_onnx_file_path INPUT_ONNX_FILE_PATH
      Input onnx file path.
  
    -b BATCH_SIZE, --batch_size BATCH_SIZE
      Value to be substituted if input batch size is undefined.
      This is ignored if the input dimensions are all of static size.
      Also ignored if input_numpy_file_paths_for_testing or
      numpy_ndarrays_for_testing or fixed_shapes is specified.
  
    -fs FIXED_SHAPES [FIXED_SHAPES ...], --fixed_shapes FIXED_SHAPES [FIXED_SHAPES ...]
      Input OPs with undefined shapes are changed to the specified shape.
      This parameter can be specified multiple times depending on
      the number of input OPs in the model.
      Also ignored if input_numpy_file_paths_for_testing or
      numpy_ndarrays_for_testing is specified.
      e.g.
      --fixed_shapes 1 3 224 224 \
      --fixed_shapes 1 5 \
      --fixed_shapes 1 1 224 224
  
    -tlc TEST_LOOP_COUNT, --test_loop_count TEST_LOOP_COUNT
      Number of times to run the test.
      The total execution time is divided by the number of times the test is executed,
      and the average inference time per inference is displayed.
  
    -oep {tensorrt,cuda,openvino_cpu,openvino_gpu,cpu}, \
      --onnx_execution_provider {tensorrt,cuda,openvino_cpu,openvino_gpu,cpu}
  
      ONNX Execution Provider.
  
    -ifp INPUT_NUMPY_FILE_PATHS_FOR_TESTING, \
      --input_numpy_file_paths_for_testing INPUT_NUMPY_FILE_PATHS_FOR_TESTING
  
      Use an external file of numpy.ndarray saved using np.save as input data for testing.
      This parameter can be specified multiple times depending on
      the number of input OPs in the model.
      If this parameter is specified, the value specified for
      batch_size and fixed_shapes are ignored.
      e.g.
      --input_numpy_file_paths_for_testing aaa.npy \
      --input_numpy_file_paths_for_testing bbb.npy \
      --input_numpy_file_paths_for_testing ccc.npy
  
    -ofp, --output_numpy_file
      Outputs the last inference result to an .npy file.
  
    -n, --non_verbose
      Do not show all information logs. Only error logs are displayed.

• Pythonスクリプトから利用する場合

  >>> from sit4onnx import inference
  >>> help(inference)
  
  Help on function inference in module sit4onnx.onnx_inference_test:
  
  inference(
    input_onnx_file_path: str,
    batch_size: Union[int, NoneType] = 1,
    fixed_shapes: Union[List[int], NoneType] = None,
    test_loop_count: Union[int, NoneType] = 10,
    onnx_execution_provider: Union[str, NoneType] = 'tensorrt',
    input_numpy_file_paths_for_testing: Union[List[str], NoneType] = None,
    numpy_ndarrays_for_testing: Union[List[numpy.ndarray], NoneType] = None,
    output_numpy_file: Union[bool, NoneType] = False,
    non_verbose: Union[bool, NoneType] = False
  ) -> List[numpy.ndarray]
  
      Parameters
      ----------
      input_onnx_file_path: str
          Input onnx file path.
  
      batch_size: Optional[int]
          Value to be substituted if input batch size is undefined.
          This is ignored if the input dimensions are all of static size.
          Also ignored if input_numpy_file_paths_for_testing or
          numpy_ndarrays_for_testing is specified.
          Default: 1
  
      fixed_shapes: Optional[List[int]]
          Input OPs with undefined shapes are changed to the specified shape.
          This parameter can be specified multiple times depending on the number of input OPs
          in the model.
          Also ignored if input_numpy_file_paths_for_testing or numpy_ndarrays_for_testing
          is specified.
          e.g.
              [
                  [1, 3, 224, 224],
                  [1, 5],
                  [1, 1, 224, 224],
              ]
          Default: None
  
      test_loop_count: Optional[int]
          Number of times to run the test.
          The total execution time is divided by the number of times the test is executed,
          and the average inference time per inference is displayed.
          Default: 10
  
      onnx_execution_provider: Optional[str]
          ONNX Execution Provider.
          "tensorrt" or "cuda" or "openvino_cpu" or "openvino_gpu" or "cpu"
          Default: "tensorrt"
  
      input_numpy_file_paths_for_testing: Optional[List[str]]
          Use an external file of numpy.ndarray saved using np.save as input data for testing.
          If this parameter is specified, the value specified for batch_size and fixed_shapes
          are ignored.
          numpy_ndarray_for_testing Cannot be specified at the same time.
          For models with multiple input OPs, specify multiple numpy file paths in list format.
          e.g. ['aaa.npy', 'bbb.npy', 'ccc.npy']
          Default: None
  
      numpy_ndarrays_for_testing: Optional[List[np.ndarray]]
          Specify the numpy.ndarray to be used for inference testing.
          If this parameter is specified, the value specified for batch_size and fixed_shapes
          are ignored.
          input_numpy_file_paths_for_testing Cannot be specified at the same time.
          For models with multiple input OPs, specify multiple numpy.ndarrays in list format.
          e.g.
          [
              np.asarray([[[1.0],[2.0],[3.0]]], dtype=np.float32),
              np.asarray([1], dtype=np.int64),
          ]
          Default: None
  
      output_numpy_file: Optional[bool]
          Outputs the last inference result to an .npy file.
          Default: False
  
      non_verbose: Optional[bool]
          Do not show all information logs. Only error logs are displayed.
          Default: False
  
      Returns
      -------
      final_results: List[np.ndarray]
          Last Reasoning Results.

起動オプションが豊富過ぎて少々混乱するかもしれませんが、手軽にONNXの動作を確認するために作成したツールですので最低限指定すべきオプションは -if と -oep の２つのみです。-if で動作確認したいONNXファイルを指定し、-oep で ランタイム(エグゼキューションプロバイダ) を選びます。-oep cpu という表記はCPU上でONNX推論を行うことを表し、TensorRTで推論したいときは -oep tensorrt、CUDAで推論したいときは -oep cuda と指定します。CPU推論で動作確認をするときのコマンドは下記のとおりです。

sit4onnx -if LayerNormalization2D_11.onnx -oep cpu

実行結果は下記の通りです。sit4onnx の内部でONNXのモデル構造を読み取り、入力OPの形状から自動的にダミーのテスト用テンソルを生成してパフォーマンスを計測します。デフォルトは内部で11回ループし、初回ループの計測結果を捨てて、残りの10回分の推論パフォーマンスを計測しています。初回ループの計測結果を捨てているのは、初回モデルのロード時のみ前処理のオーバーヘッドの影響を大きく受けてパフォーマンスが著しく悪化するため、前段のオーバーヘッド部分を除いた推論のパフォーマンスのみを抽出しています。一部界隈では初回オーバーヘッドのことをウォームアップと呼んでいるようです。ベンチマーク結果の単位時間はミリ秒です。

INFO: file: LayerNormalization2D_11.onnx
INFO: providers: ['CPUExecutionProvider']
INFO: input_name.1: LayerNormalization2D_input shape: [20, 5, 10, 10] dtype: float32
INFO: test_loop_count: 10
INFO: total elapsed time:  0.3662109375 ms
INFO: avg elapsed time per pred:  0.03662109375 ms
INFO: output_name.1: LayerNormalization2D_output shape: [20, 5, 10, 10] dtype: float32

コンソールの表示イメージです。

ちなみに、TensorRT や CUDA で推論のベンチマークを行いたい場合は、３章で導入した各種パッケージのうち、onnxruntime あるいは onnxruntime-gpu をTensorRTバックエンドを有効な状態にしてカスタムビルドしたパッケージに差し替える必要があります。ご自身の環境に合う TensorRT対応 onnxruntime-gpu のインストーラ生成の方法は下記のとおりです。

注意点として、TensorRT 8.4.0 より新しいバージョンが導入された環境の場合、FP16の演算精度が著しく劣化することがある問題を確認しています。従って、私のHostPC上のTensorRTバージョンは 8.4.0 EA という少し古いバージョンを導入しています。

• Ubuntu 20.04上でのインストーラ生成サンプルコマンド

  git clone https://github.com/microsoft/onnxruntime.git \
  && cd onnxruntime
  git checkout 49d7050b88338dd57839159aa4ce8fb0c199b064
  
  dpkg -l | grep TensorRT
  
  ii graphsurgeon-tf        8.4.0-1+cuda11.6   amd64 GraphSurgeon for TensorRT package
  ii libnvinfer-bin         8.4.0-1+cuda11.6   amd64 TensorRT binaries
  ii libnvinfer-dev         8.4.0-1+cuda11.6   amd64 TensorRT development libraries and headers
  ii libnvinfer-doc         8.4.0-1+cuda11.6   all   TensorRT documentation
  ii libnvinfer-plugin-dev  8.4.0-1+cuda11.6   amd64 TensorRT plugin libraries
  ii libnvinfer-plugin8     8.4.0-1+cuda11.6   amd64 TensorRT plugin libraries
  ii libnvinfer-samples     8.4.0-1+cuda11.6   all   TensorRT samples
  ii libnvinfer8            8.4.0-1+cuda11.6   amd64 TensorRT runtime libraries
  ii libnvonnxparsers-dev   8.4.0-1+cuda11.6   amd64 TensorRT ONNX libraries
  ii libnvonnxparsers8      8.4.0-1+cuda11.6   amd64 TensorRT ONNX libraries
  ii libnvparsers-dev       8.4.0-1+cuda11.6   amd64 TensorRT parsers libraries
  ii libnvparsers8          8.4.0-1+cuda11.6   amd64 TensorRT parsers libraries
  ii onnx-graphsurgeon      8.4.0-1+cuda11.6   amd64 ONNX GraphSurgeon for TensorRT package
  ii python3-libnvinfer     8.4.0-1+cuda11.6   amd64 Python 3 bindings for TensorRT
  ii python3-libnvinfer-dev 8.4.0-1+cuda11.6   amd64 Python 3 development package for TensorRT
  ii tensorrt               8.4.0.6-1+cuda11.6 amd64 Meta package of TensorRT
  ii uff-converter-tf       8.4.0-1+cuda11.6   amd64 UFF converter for TensorRT package
  
  sudo chmod +x build.sh
  sudo pip install cmake --upgrade
  
  ./build.sh \
  --config Release \
  --cudnn_home /usr/lib/x86_64-linux-gnu/ \
  --cuda_home /usr/local/cuda \
  --use_tensorrt \
  --use_cuda \
  --tensorrt_home /usr/src/tensorrt/ \
  --enable_pybind \
  --build_wheel \
  --parallel $(nproc) \
  --skip_tests
  
  find . -name "*.whl"
  ./build/Linux/Release/dist/onnxruntime_gpu-1.12.0-cp38-cp38-linux_x86_64.whl
  
  pip install ./build/Linux/Release/dist/onnxruntime_gpu-*.whl

  sit4onnx -if LayerNormalization2D_11.onnx -oep tensorrt

あるいは、TensorRT やその他の主要MLフレームワークを全てカスタムビルドして導入済みの巨大Docker Imageをダウンロードして利用する方法もあります。23GBほどあります。

• 主要MLフレームワーク全部載せDockerコンテナの起動サンプル

  docker pull ghcr.io/pinto0309/openvino2tensorflow:latest
  
  docker run --gpus all -it --rm \
  -v `pwd`:/home/user/workdir \
  ghcr.io/pinto0309/openvino2tensorflow:latest

  sit4onnx -if LayerNormalization2D_11.onnx -oep tensorrt

TensorRTを使用したベンチマーク結果は下記のとおりです。ONNXファイルを TRT Engine という形式にコンパイルする動作が実行されるため、ベンチマークがスタートするまでに１分ほど待つ必要があります。

INFO: file: LayerNormalization2D_11.onnx
INFO: providers: ['TensorrtExecutionProvider', 'CPUExecutionProvider']
INFO: input_name.1: LayerNormalization2D_input shape: [20, 5, 10, 10] dtype: float32
INFO: test_loop_count: 10
INFO: total elapsed time:  0.7319450378417969 ms
INFO: avg elapsed time per pred:  0.07319450378417969 ms
INFO: output_name.1: LayerNormalization2D_output shape: [20, 5, 10, 10] dtype: float32

演算精度はFP16ですが、CPUのほうが２倍ほど速いようです。
CUDAではどうでしょうか。

sit4onnx -if LayerNormalization2D_11.onnx -oep cuda

INFO: file: LayerNormalization2D_11.onnx
INFO: providers: ['CUDAExecutionProvider', 'CPUExecutionProvider']
INFO: input_name.1: LayerNormalization2D_input shape: [20, 5, 10, 10] dtype: float32
INFO: test_loop_count: 10
INFO: total elapsed time:  0.7827281951904297 ms
INFO: avg elapsed time per pred:  0.07827281951904297 ms
INFO: output_name.1: LayerNormalization2D_output shape: [20, 5, 10, 10] dtype: float32

ほとんど性能差は出ませんでした。
比較がしにくいため、CPUとTensorRTとCUDAのベンチマーク結果を並べてみます。

• CPU – Intel® Core i9-10900K CPU @ 3.70GHz 最大20スレッド

  INFO: file: LayerNormalization2D_11.onnx
  INFO: providers: ['CPUExecutionProvider']
  INFO: input_name.1: LayerNormalization2D_input shape: [20, 5, 10, 10] dtype: float32
  INFO: test_loop_count: 10
  INFO: total elapsed time:  0.3662109375 ms
  INFO: avg elapsed time per pred:  0.03662109375 ms
  INFO: output_name.1: LayerNormalization2D_output shape: [20, 5, 10, 10] dtype: float32

• TensorRT 8.4.0 + CUDA 11.6 – RTX3070

  INFO: file: LayerNormalization2D_11.onnx
  INFO: providers: ['TensorrtExecutionProvider', 'CPUExecutionProvider']
  INFO: input_name.1: LayerNormalization2D_input shape: [20, 5, 10, 10] dtype: float32
  INFO: test_loop_count: 10
  INFO: total elapsed time:  0.7319450378417969 ms
  INFO: avg elapsed time per pred:  0.07319450378417969 ms
  INFO: output_name.1: LayerNormalization2D_output shape: [20, 5, 10, 10] dtype: float32

• CUDA 11.6 – RTX3070

  INFO: file: LayerNormalization2D_11.onnx
  INFO: providers: ['CUDAExecutionProvider', 'CPUExecutionProvider']
  INFO: input_name.1: LayerNormalization2D_input shape: [20, 5, 10, 10] dtype: float32
  INFO: test_loop_count: 10 INFO: total elapsed time:  0.7827281951904297 ms
  INFO: avg elapsed time per pred:  0.07827281951904297 ms
  INFO: output_name.1: LayerNormalization2D_output shape: [20, 5, 10, 10] dtype: float32

プリミティブな演算のみで構成された opset=11 の LayerNormalization を演算する場合は、私のHostPC上ではCPUのほうが速いようです。ただし、今回はFP16の精度のみでTensorRTの演算速度を検証したため、Int8へ量子化した場合は最大で 2.5倍 ほど速くなる可能性はあります。

念の為、opset=17 で生成した LayerNormalization でも比較しておきます。

• CPU – Intel® Core i9-10900K CPU @ 3.70GHz 最大20スレッド

  INFO: file: LayerNormalization2D_17.onnx
  INFO: providers: ['CPUExecutionProvider']
  INFO: input_name.1: LayerNormalization2D_input shape: [20, 5, 10, 10] dtype: float32
  INFO: test_loop_count: 10
  INFO: total elapsed time:  0.4291534423828125 ms
  INFO: avg elapsed time per pred:  0.04291534423828125 ms
  INFO: output_name.1: LayerNormalization2D_output shape: [20, 5, 10, 10] dtype: float32

• TensorRT 8.4.0 + CUDA 11.6 – RTX3070

  2022-11-10 06:55:15.182974909 [W:onnxruntime:Default, tensorrt_execution_provider.cc:1046 GetCapability] [TensorRT EP] No graph will run on TensorRT execution provider
  2022-11-10 06:55:15.949520352 [W:onnxruntime:Default, tensorrt_execution_provider.h:60 log] [2022-11-10 06:55:15 WARNING] external/onnx-tensorrt/onnx2trt_utils.cpp:367: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
  INFO: file: LayerNormalization2D_17.onnx
  INFO: providers: ['TensorrtExecutionProvider', 'CPUExecutionProvider']
  INFO: input_name.1: LayerNormalization2D_input shape: [20, 5, 10, 10] dtype: float32
  EP Error: [ONNXRuntimeError] : 11 : EP_FAIL : Non-zero status code returned while running TRTKernel_graph_torch_jit_2590721683996607622_1 node. Name:'TensorrtExecutionProvider_TRTKernel_graph_torch_jit_2590721683996607622_1_0' Status Message: TensorRT EP output tensor data type: 0 not supported. using ['TensorrtExecutionProvider', 'CPUExecutionProvider']
  Falling back to ['CUDAExecutionProvider', 'CPUExecutionProvider'] and retrying.
  INFO: test_loop_count: 10
  INFO: total elapsed time:  1.1851787567138672 ms
  INFO: avg elapsed time per pred:  0.11851787567138672 ms
  INFO: output_name.1: LayerNormalization2D_output shape: [20, 5, 10, 10] dtype: float32

• CUDA 11.6 – RTX3070

  INFO: file: LayerNormalization2D_17.onnx
  INFO: providers: ['CUDAExecutionProvider', 'CPUExecutionProvider']
  INFO: input_name.1: LayerNormalization2D_input shape: [20, 5, 10, 10] dtype: float32
  INFO: test_loop_count: 10
  INFO: total elapsed time:  1.1889934539794922 ms
  INFO: avg elapsed time per pred:  0.11889934539794922 ms
  INFO: output_name.1: LayerNormalization2D_output shape: [20, 5, 10, 10] dtype: float32

CPU推論が遅くなりましたね。TensorRT 8.4.0 に至っては EP Error が発生し、結果として CUDA にフォールバックしてしまいました。CUDAもCPUとは比べ物にならないほど遅いようです。特定のTensorRTのバージョンとCUDAのバージョンの組み合わせではオペレーションの対応状況が異なることが想定されますので上記のベンチマーク結果が全てではありませんし、最新のバージョンの組み合わせであれば期待通りのベンチマークが出る可能性があります。

さて、このベンチマークはほとんど意味のない比較であることをご理解願います。LayerNormalization 単体で構成されたモデルなど実用的なモデルではありません。そもそも真面目に推論速度を比較することがこの章でご紹介したかったことの本質ではありません。あくまで、簡易的にモデルを生成し、簡易的に推論のベンチマークを行う手段をご紹介したに過ぎません。ベンチマークプログラムを動かすにあたってどこがスループットのボトルネックになっているかまでは一切未検証です。もしかしたら、GPUへのテンソルの転送コストがネックになって全体の処理パフォーマンスが遅いように見えているだけかもしれません。従いまして、実用的な構成のモデルでベンチマークをとってみることのほうが重要です。

5. モデル生成の基礎 [ONNXのみ]

この章では、ONNXのモデル構造の基礎、PyTorchを経由せずにONNXのオペレーションおよびモデル構造を生成する方法をご紹介します。

5-1. ONNXモデルの構造

モデルの動作を微調整するためだけにPyTorchをインストールした環境を用意したくありません。CUDAのバージョンアンマッチや古いバージョンのPyTorchを動作させる環境を整えることにとても手間が掛かることがあります。ただモデルの微調整をする用途のためだけに用意するPyTorchは機能過剰です。すでに手元にあるONNXファイルを少し調整してフローを変えたり、演算の一部を別の演算に置き換えたりしたいときがあります。世界各国のリサーチャーの趣味嗜好が入り混じったPyTorchのソースコードをできるだけ読みたくないのです。他の方々が書いたソースコードを解析するのは非常にコストが掛かります。皆さんがどれほどそういった感覚を持たれるかどうかは分かりませんが、ちょっとした調整をする場合は構造的に綺麗に整理された状態のONNXファイルをNetronで眺めながら課題となる箇所を特定するほうがはるかに手間が少ないときがあります。

したがって、PyTorchに一切頼らずにONNXを直接生成することを考えますが、直接生成するためにはONNXの構造をまず知る必要があります。構造に関しては公式リポジトリにとても詳細な説明があります。Extensible computation graph model 情報量がとても多く理解するのがとても大変ですので、ONNXモデルを生成したりチューニングする際によく注目するポイントなどを絞ってご紹介します。

では、前章で生成した LayerNormalization をもとにして構造を解析してみます。使用するモデルは opset=17 で生成した LayerNormalization です。

• モデル全体の定義

  

• LayerNormalization の定義

  

まずは我々人間が構造を把握しやすいようにONNXファイルをJSONに変換してみます。ただ、JSON化するためにわざわざ変換プログラムを書きたくありませんので、自作のツール onnx2json を使用します。

• コマンドラインインタフェース

  onnx2json -h
  
  usage:
    onnx2json [-h]
    -if INPUT_ONNX_FILE_PATH
    -oj OUTPUT_JSON_PATH
    [-i JSON_INDENT]
  
  optional arguments:
    -h, --help
        show this help message and exit
  
    -if INPUT_ONNX_FILE_PATH, --input_onnx_file_path INPUT_ONNX_FILE_PATH
        Input ONNX model path. (*.onnx)
  
    -oj OUTPUT_JSON_PATH, --output_json_path OUTPUT_JSON_PATH
        Output JSON file path (*.json) If not specified, no JSON file is output.
  
    -i JSON_INDENT, --json_indent JSON_INDENT
        Number of indentations in JSON. (default=2)

• Pythonスクリプトから利用する場合

  >>> from onnx2json import convert
  >>> help(convert)
  
  Help on function convert in module onnx2json.onnx2json:
  
  convert(
    input_onnx_file_path: Union[str, NoneType] = '',
    onnx_graph: Union[onnx.onnx_ml_pb2.ModelProto, NoneType] = None,
    output_json_path: Union[str, NoneType] = '',
    json_indent: Union[int, NoneType] = 2
  )
  
      Parameters
      ----------
      input_onnx_file_path: Optional[str]
          Input onnx file path.
          Either input_onnx_file_path or onnx_graph must be specified.
          Default: ''
  
      onnx_graph: Optional[onnx.ModelProto]
          onnx.ModelProto.
          Either input_onnx_file_path or onnx_graph must be specified.
          onnx_graph If specified, ignore input_onnx_file_path and process onnx_graph.
  
      output_onnx_file_path: Optional[str]
          Output onnx file path. If not specified, no ONNX file is output.
          Default: ''
  
      json_indent: Optional[int]
          Number of indentations in JSON.
          Default: 2
  
      Returns
      -------
      onnx_json: dict
          Converted JSON dict.

３章の環境構築を正しく行っている場合はすでにインストール済みです。ONNXをJSONへ変換するコマンドは下記のとおりです。

onnx2json -if LayerNormalization2D_17.onnx -oj LayerNormalization2D_17.json

生成されたJSONファイルをVSCodeなどのIDEで開いて内容を確認してみます。rawData の部分のみ省略してJSON全体を転記したものが下記です。特に説明をしなくてもなんとなく雰囲気がつかめるレベルのシンプルさだと思います。直接書き換えてONNXファイルへ戻すことも可能です。

{
  "irVersion": "8",
  "producerName": "pytorch",
  "producerVersion": "1.14.0",
  "graph": {
    "node": [
      {
        "input": [
          "LayerNormalization2D_input",
          "/Constant_output_0",
          "/Constant_1_output_0"
        ],
        "output": [
          "LayerNormalization2D_output"
        ],
        "name": "/LayerNormalization",
        "opType": "LayerNormalization",
        "attribute": [
          {
            "name": "axis",
            "i": "-3",
            "type": "INT"
          },
          {
            "name": "epsilon",
            "f": 1e-05,
            "type": "FLOAT"
          }
        ]
      }
    ],
    "name": "torch_jit",
    "initializer": [
      {
        "dims": [
          "5",
          "10",
          "10"
        ],
        "dataType": 1,
        "name": "/Constant_output_0",
        "rawData": "zczMPc3MzD3NzM...Mw9zczMPc3MzD0="
      },
      {
        "dims": [
          "5",
          "10",
          "10"
        ],
        "dataType": 1,
        "name": "/Constant_1_output_0",
        "rawData": "zcxMPs3MTDNzEw...zw+zcxMPs3MTD4="
      }
    ],
    "input": [
      {
        "name": "LayerNormalization2D_input",
        "type": {
          "tensorType": {
            "elemType": 1,
            "shape": {
              "dim": [
                {
                  "dimValue": "20"
                },
                {
                  "dimValue": "5"
                },
                {
                  "dimValue": "10"
                },
                {
                  "dimValue": "10"
                }
              ]
            }
          }
        }
      }
    ],
    "output": [
      {
        "name": "LayerNormalization2D_output",
        "type": {
          "tensorType": {
            "elemType": 1,
            "shape": {
              "dim": [
                {
                  "dimValue": "20"
                },
                {
                  "dimValue": "5"
                },
                {
                  "dimValue": "10"
                },
                {
                  "dimValue": "10"
                }
              ]
            }
          }
        }
      }
    ]
  },
  "opsetImport": [
    {
      "domain": "",
      "version": "17"
    }
  ]
}

• 各値の意味 – 公式の仕様に書かれていない情報も追記

  Key	Value
  irVersion	ONNXのIRフォーマットのバージョン番号
  producerName	モデルを生成するために使用したツールの名前
  producerVersion	モデルを生成するために使用したツールのバージョン
  graph	モデルのグラフ本体の定義 Key Value  node オペレーションの入出力の接続関係をリスト形式で表現  name グラフの名前  initializer オペレーションに含まれる定数、Convの重みやテンソルの形状 Key Value  dims 定数の形状、Numpyのndarrayの各次元のサイズ  dataType 定数のデータ型、inputのelemTypeを参照  name 定数の名前  rawData Float32やFloat16、Int8などの各種精度の Numpy.ndarrayがBase64 エンコードされた値 Float32:[1,2,3,4] AACAPwAAAEAAAEBAAACAQA== Int64:[1,2] AQAAAAAAAAACAAAAAAAAAA==  input 入力変数、一部省略 Key Value  elemType 変数のデータ型 Type Values Type Values  float16 10 int8 3  float32 1 int16 5  float64 11 int32 6  bool 9 int64 7  uint8 2    uint16 4  uint32 12  uint64 13  dim dimValue: 数値、固定形状 dimParam: 文字列、可変形状  output 出力変数、’input’と同じ
  opsetImport	opsetのドメイン名とopsetのバージョン番号 Key Value  domain ドメイン名  version バージョン番号 ドメイン名は ai.onnx vXX com.microsoft v1 com.microsoft.nchwc v1 ai.onnx.training v1 ai.onnx.preview.training v1 com.microsoft.experimental v1 など

ONNXのJSONを手で書き換えて保存し、再びONNXファイルへ戻すテクニックは後続の章で取り上げます。ちなみに、Numpy.ndarrayをBase64エンコードした値を知りたい場合は sed4onnx を使用すると簡単です。

sed4onnx --constant_string [1,2] --dtype int64 --mode encode
AQAAAAAAAAACAAAAAAAAAA==

sed4onnx --constant_string AQAAAAAAAAACAAAAAAAAAA== --dtype int64 --mode decode
[1,2]

畳み込みの重みやバイアス、Normalizationのバイアスなどは全てこのルールに則って参照・書き換えが可能です。

なお捕捉ですが、dimValue と dimParam の用法は下記のイメージです。例えば、バッチサイズが可変の定義 N のときは "dimParam: "N" となります。dimParam が指定された形状部分は未定義の次元として任意の文字列を指定することができます。サンプルは N を例示しましたが、batch_size でも良いですし、abcdefg でも -1 でも問題ありません。指定された値は全てただの表示用の文字列として扱われる仕様です。dimValue を使用する場合は正の整数のみ指定可能です。従って、-1 は文字列扱いとなりますので指定できません。

"input": [
      {
        "name": "Abs_input",
        "type": {
          "tensorType": {
            "elemType": 1,
            "shape": {
              "dim": [
                {
                  "dimParam": "N"
                },
                {
                  "dimValue": "100"
                }
              ]
            }
          }
        }
      }
    ],

実は、ココまでの基本が押さえられていればJSONをIDEで手書きしてONNXモデルをフルスクラッチで生成することが可能です。ただ、現実的にはJSONをゼロから書き上げるのは辛いですので、すでに手元にあるONNXファイルを微改造する程度のシチュエーションのほうが多いかもしれません。また運用上、モデルのチューニングを自動化するなどの要件がある場合は手書きはナンセンスですので、後続の章でご紹介するようなモデル加工用のツールをシェルで連続的に実行して自動化するのが現実的です。

5-2. onnx-graphsurgeon

ONNXモデルの構造を把握できたところで、この章以降で多用するとても重要で便利なツールをご紹介しておきます。NVIDIAさんが公開してくれている onnx-graphsurgeon というツールです。TensorRTの附属ツールとして公開されているため少し見つけ出しにくい場所にリソースが存在していますが、nvidia-pyindex というパッケージインデックスを介してpipパッケージを取得することができます。３章の手順を実行している場合はすでに導入済みになっています。surgeon は 外科医 と日本語訳され、その名のイメージの通りONNXのモデル(グラフ)を切除したり結合したりと自由自在に加工することができます。モデル加工のチュートリアルが数種類コミットされていますので、ご興味が有る方は examples フォルダを覗いてみると良いと思います。ただ、このツールの仕様や使い方をマスターすることがこのブログの本質ではないため、後続の章でご紹介するテクニックではこのツールを皆さんが直接呼び出すことはありません。全て自作ツールでラップして分かりにくい仕様部分をブラックボックス化した状態で利用します。

便利ですが、Pythonでソースコードを書いて onnx_graphsurgeon を適切に呼び出してモデルを生成したり加工する必要があります。ツールの仕様を理解することにコストを掛けたくないのが人情ですが、理解しておくと自作ツールを作成したりよりトリッキーな加工を行いたい要件が出てきたときに役に立つ時があるかもしれません。従いまして、onnx_graphsurgeon の examples の中のチュートリアルのひとつだけをココではご紹介します。

• 01_creating_a_model

  import onnx_graphsurgeon as gs
  import numpy as np
  import onnx
  
  OPSET=11
  X = gs.Variable(name="X", dtype=np.float32, shape=(1, 3, 5, 5))
  Y = gs.Variable(name="Y", dtype=np.float32, shape=(1, 3, 1, 1))
  node = gs.Node(op="GlobalLpPool", attrs={"p": 2}, inputs=[X], outputs=[Y])
  
  graph = gs.Graph(nodes=[node], inputs=[X], outputs=[Y], opset=OPSET)
  onnx.save(gs.export_onnx(graph), f"GlobalLpPool_{OPSET}.onnx")

上記のソースコードを任意の .py ファイルとして保存して実行すると、下図のようなモデルが生成されます。

READMEに該当するドキュメントにはほとんど説明が書かれていないため、初見では意味が分からないことが多いと思います。onnx_graphsurgeon の構成要素をまずはご紹介します。

Pythonのコードは小さな変数 Variable や定数 Constant から定義を初めて Variable を介して Node 間を結合していき、Graph を生成するイメージです。オペレーションごとにONNXの仕様に則って全ての変数や定数をあらかじめ定義しておく必要があり少々直感的ではありませんが、簡単なコードを書くことでモデルを構築できる雰囲気はご理解いただけると思います。

onnx.save(gs.export_onnx(graph), "test_globallppool.onnx") の部分で onnx_graphsurgeon.Graph 形式のオブジェクトをONNXファイルとして出力しています。gs.export_onnx(graph) の部分は onnx.ModelProto という、ONNXファイルへ出力可能な Protocol Buffers 形式のオブジェクトを生成してくれています。onnx.save(...) はONNXの標準APIの save メソッドに onnx.ModelProto を引き渡しているだけです。

ここで繰り返しになりますが、便利で柔軟な対応が可能なツールではありますが、このツールの仕様を覚えることが本質ではありません。

5-3. Node生成

前述の onnx_graphsurgeon を使用してPythonのソースコードを書いて実行すれば簡単にONNXのモデルを生成できることはご理解いただけたかと思います。しかし、そもそもPythonのソースコードを書きたくないです。できればシェルで期待通りのオペレーションをポンッと生成したいです。細々としたソースコードを書く数十分の時間すら惜しいです。そこで、onnx_graphsurgeon をラップしてシェル操作だけでONNXのノードを生成できる sog4onnx というツールを作成しましたのでそちらを使用してNode生成を行ってみます。

Nodeを生成するために把握しなければいけないのはONNXのオペレーションの仕様です。GlobalLpPool の仕様はこちらです。GlobalLpPool 読んでいただくと分かるとおり、特に難しいことをしているオペレーションでは無さそうです。入力にひとつの値と属性値 p をとり、出力はひとつの値となるようです。属性値Attributes p は Int64、入力Variable X は Float16 又は Float32 又は Float64、出力Variable Y は Float16 又は Float32 又は Float64 の型のみが許容されるようです。

• GlobalLpPool の仕様

  GlobalLpPool-2
  
  GlobalLpPool consumes an input tensor X and applies lp pool pooling across
  the values in the same channel. This is equivalent to
  LpPool with kernel size equal to the spatial dimension of input tensor.
  
  - Version
      This version of the operator has been available since version 2
      of the default ONNX operator set.
  
  - Attributes
      p : int (default is 2)
          p value of the Lp norm used to pool over the input data.
  
  - Inputs
      X (differentiable) : T
          Input data tensor from the previous operator; dimensions
          for image case are (N x C x H x W), where N is the batch size,
          C is the number of channels, and H and W are the height and
          the width of the data. For non image case, the dimensions are
          in the form of (N x C x D1 x D2 ... Dn), where N is the batch size.
  
  - Outputs
      Y (differentiable) : T
          Output data tensor from pooling across the input tensor.
          The output tensor has the same rank as the input.
          The first two dimensions of output shape are the same
          as the input (N x C), while the other dimensions are all 1.
  
  - Type Constraints
      T : tensor(float16), tensor(float), tensor(double)
          Constrain input and output types to float tensors.

では前節で onnx_graphsurgeon を使用して生成したオペレーションと同じ内容のオペレーションを sog4onnx を使用して生成してみます。

sog4onnx \
--op_type GlobalLpPool \
--opset 11 \
--op_name GlobalLpPool1 \
--input_variables X float32 [1,3,5,5] \
--output_variables Y float32 [1,3,1,1] \
--attributes p int64 2 \
--output_onnx_file_path GlobalLpPool_11_sog.onnx

生成できました。onnx_graphsurgeon で生成したモデルと同じです。バックエンドで onnx_graphsurgeon をコールしていることには変わりありませんが記述量が大幅に減りました。より直感的にONNXのオペレーションを生成できます。またシェルスクリプトへ記述して機械的に、連続的にオペレーションを生成することができます。

念の為ツールの仕様を記載します。op_type、opset、op_name、input_variables、output_variables、attributes そして output_onnx_file_path を必要数分だけ順番に指定していくだけです。

• コマンドラインインタフェース

  sog4onnx -h
  
  usage: sog4onnx [-h]
    --ot OP_TYPE
    --os OPSET
    --on OP_NAME
    [-iv NAME TYPE VALUE]
    [-ov NAME TYPE VALUE]
    [-a NAME DTYPE VALUE]
    [-of OUTPUT_ONNX_FILE_PATH]
    [-n]
  
  optional arguments:
    -h, --help
      show this help message and exit
  
    -ot OP_TYPE, --op_type OP_TYPE
      ONNX OP type.
      https://github.com/onnx/onnx/blob/main/docs/Operators.md
  
    -os OPSET, --opset OPSET
      ONNX opset number.
  
    -on OP_NAME, --op_name OP_NAME
      OP name.
  
    -iv INPUT_VARIABLES INPUT_VARIABLES INPUT_VARIABLES,
          --input_variables INPUT_VARIABLES INPUT_VARIABLES INPUT_VARIABLES
      input_variables can be specified multiple times.
      --input_variables variable_name numpy.dtype shape
      https://github.com/onnx/onnx/blob/main/docs/Operators.md
  
      e.g.
      --input_variables i1 float32 [1,3,5,5] \
      --input_variables i2 int32 [1] \
      --input_variables i3 float64 [1,3,224,224]
  
    -ov OUTPUT_VARIABLES OUTPUT_VARIABLES OUTPUT_VARIABLES,
          --output_variables OUTPUT_VARIABLES OUTPUT_VARIABLES OUTPUT_VARIABLES
      output_variables can be specified multiple times.
      --output_variables variable_name numpy.dtype shape
      https://github.com/onnx/onnx/blob/main/docs/Operators.md
  
      e.g.
      --output_variables o1 float32 [1,3,5,5] \
      --output_variables o2 int32 [1] \
      --output_variables o3 float64 [1,3,224,224]
  
    -a ATTRIBUTES ATTRIBUTES ATTRIBUTES, --attributes ATTRIBUTES ATTRIBUTES ATTRIBUTES
      attributes can be specified multiple times.
      dtype is one of "float32" or "float64" or "int32" or "int64" or "str".
      --attributes name dtype value
      https://github.com/onnx/onnx/blob/main/docs/Operators.md
  
      e.g.
      --attributes alpha float32 1.0 \
      --attributes beta float32 1.0 \
      --attributes transA int32 0 \
      --attributes transB int32 0
  
    -of OUTPUT_ONNX_FILE_PATH, --output_onnx_file_path OUTPUT_ONNX_FILE_PATH
      Output onnx file path.
      If not specified, a file with the OP type name is generated.
  
      e.g. op_type="Gemm" -> Gemm.onnx
  
    -n, --non_verbose
      Do not show all information logs. Only error logs are displayed.

• Pythonスクリプトから利用する場合

  >> from sog4onnx import generate
  >>> help(generate)
  Help on function generate in module sog4onnx.onnx_operation_generator:
  
  generate(
    op_type: str,
    opset: int,
    op_name: str,
    input_variables: dict,
    output_variables: dict,
    attributes: Union[dict, NoneType] = None,
    output_onnx_file_path: Union[str, NoneType] = '',
    non_verbose: Union[bool, NoneType] = False
  ) -> onnx.onnx_ml_pb2.ModelProto
  
      Parameters
      ----------
      op_type: str
          ONNX op type.
          See below for the types of OPs that can be specified.
          https://github.com/onnx/onnx/blob/main/docs/Operators.md
  
          e.g. "Add", "Div", "Gemm", ...
  
      opset: int
          ONNX opset number.
  
          e.g. 11
  
      op_name: str
          OP name.
  
      input_variables: Optional[dict]
          Specify input variables for the OP to be generated.
          See below for the variables that can be specified.
          https://github.com/onnx/onnx/blob/main/docs/Operators.md
          {"input_var_name1": [numpy.dtype, shape], "input_var_name2": [dtype, shape], ...}
  
          e.g.
          input_variables = {
            "name1": [np.float32, [1,224,224,3]],
            "name2": [np.bool_, [0]],
            ...
          }
  
      output_variables: Optional[dict]
          Specify output variables for the OP to be generated.
          See below for the variables that can be specified.
          https://github.com/onnx/onnx/blob/main/docs/Operators.md
          {"output_var_name1": [numpy.dtype, shape], "output_var_name2": [dtype, shape], ...}
  
          e.g.
          output_variables = {
            "name1": [np.float32, [1,224,224,3]],
            "name2": [np.bool_, [0]],
            ...
          }
  
      attributes: Optional[dict]
          Specify output attributes for the OP to be generated.
          See below for the attributes that can be specified.
          When specifying Tensor format values, specify an array converted to np.ndarray.
          https://github.com/onnx/onnx/blob/main/docs/Operators.md
          {"attr_name1": value1, "attr_name2": value2, "attr_name3": value3, ...}
  
          e.g.
          attributes = {
            "alpha": 1.0,
            "beta": 1.0,
            "transA": 0,
            "transB": 0
          }
          Default: None
  
      output_onnx_file_path: Optional[str]
          Output of onnx file path.
          If not specified, no .onnx file is output.
          Default: ''
  
      non_verbose: Optional[bool]
          Do not show all information logs. Only error logs are displayed.
          Default: False
  
      Returns
      -------
      single_op_graph: onnx.ModelProto
          Single op onnx ModelProto

5-4. Nodeのコンポーネント化

前節で単体のオペレーションを生成するだけの操作に何の意味があるのか疑問を持たれた方が多いと思います。たしかにオペレーション単体のモデルでは特にメリットを感じませんが、複数のオペレーションを生成して結合し、前処理や後処理をコンポーネント化してストックしておけるとしたらどうでしょうか。PyTorchで前処理・後処理を記述することと何ら変わりはありませんが、コンポーネント化した定型的な前処理・後処理をONNXファイルとしてあらかじめ生成しておき、モデル本体部分がどのようなものに変わってもその前後の処理はあとからシェルでマージできるようになります。また、実はPyTorchで表現できないオペレーションにも対応することができるようになりますし、元のモデルがPyTorchから生成されたものではなくなおかつトレーニングコードが公開されていないモデルの改造時に役に立ちます。

コンポーネント化がどのようなことを指すのかはこちらをご覧ください。components_of_onnx READMEの下部 Progress の部分は折りたたみになっており、展開するとコンポーネント化済みのONNXがコミットされていることが分かります。

文脈として読みづらいと思いますので、一部分を抜粋してこちらに転記します。各種ハードウェア向けの特殊構造をあらかじめ実装していたり、PyTorch (TorchVision) では実装不可能なオペレーションを生成したり、定型的なノーマライゼーションの処理、RGB<->BGR変換処理を生成してコミットしています。

• ONNXの特殊コンポーネントの一部

  No.	Operator	Structure
  Z001	Normalization_rgb_imagenet
  Z002	Normalization_bgr_imagenet
  Z003	SingleClass_NonMaxSupression	480×640
  Z004	YOLACT_Edge_NonMaxSupression	550×550
  Z005	Resize_0.5×0.5	1x3xHxW->1x3x(Hx0.5)x(Wx0.5)
  Z006	Resize_HxW	1x3xHxW->1x3x(Hx?)x(Wx?)
  Z007	Myriad_workaround_NonMaxSuppression
  Z008	TensorRT_compatible_N_batch_Resize
  Z009	Unity_Barracuda_compatible_GatherND
  Z010	Unity_Barracuda_compatible_Split

では、単体のオペレーションではなく、オペレーションのセット(コンポーネント)を生成する方法をご紹介します。方法は大きく2つです。

1. PyTorchで必要なオペレーションセットの部分のみをモデル化してONNXへエクスポート
2. ONNXの単体オペレーションを複数種類生成しておいてあとからマージ

それぞれの方法にメリット・デメリットがあります。

上記の中からPyTorchを経由するとできなくてもONNXならできることを取り上げます。ONNXの特殊コンポーネントの一部 のうち、一番分かりやすいものとして Z004 の Multi-Class NonMaxSupression が該当します。と言われてもピンとこないと思いますので実例を取り上げてご説明します。

まずVision系のモデルの後処理でよく利用されるオペレーション NMS の PyTorch(TrochVision) の仕様を見てみます。Torch Vision NMS あるいは Kornia NMS ここで、Kornia の詳細には触れませんが、PyTorch向けに開発されたとても便利なメソッド群です。PyTorchでフォローできていないオペレーションが大量に提供されています。

• TorchVision NMSの仕様

  NMS
  torchvision.ops.nms(boxes: Tensor, scores: Tensor, iou_threshold: float) → Tensor[SOURCE]
  
  Performs non-maximum suppression (NMS) on the boxes according to
  their intersection-over-union (IoU).
  NMS iteratively removes lower scoring boxes which have an IoU
  greater than iou_threshold with another (higher scoring) box.
  If multiple boxes have the exact same score and satisfy
  the IoU criterion with respect to a reference box,
  the selected box is not guaranteed to be the same between CPU and GPU.
  This is similar to the behavior of argsort in PyTorch when repeated values are present.
  
  Parameters:
      boxes (Tensor[N, 4]))
          boxes to perform NMS on.
          They are expected to be in (x1, y1, x2, y2) format
          with 0 <= x1 < x2 and 0 <= y1 < y2.
      scores (Tensor[N])
          scores for each one of the boxes
      iou_threshold (float)
          discards all overlapping boxes with IoU > iou_threshold
  
  Returns:
      int64 tensor with the indices of the elements that have been kept by NMS,
      sorted in decreasing order of scores
  
  Return type:
      Tensor

何が問題なのでしょうか？ バッチ処理およびマルチクラスに対応していないことです。では、バッチ処理およびマルチクラスに対応していないとはどういうことか。boxes と scores の入力テンソルの次元がひとつ、あるいは２つ足りないです。ほとんどのシチュエーションでは不要かもしれませんが、NMSの仕様がこうなっているがゆえにとても扱いづらいシチュエーションが存在します。実際は正確に同じシチュエーションを表現している図ではありませんが下図のイメージです。NMSが80個分(80クラス分)横に展開されています。

ONNXのNMSはこうです。バッチ処理およびマルチクラス処理に対応しています。余談ですが、TensorFlow のNMSもバッチ処理に対応しています。PyTorchはTensorFlowより圧倒的に扱いやすいのは事実ですが、かゆいところに手が届きません。

NonMaxSuppression-11

Filter out boxes that have high intersection-over-union (IOU)
overlap with previously selected boxes.
Bounding boxes with score less than score_threshold are removed.
Bounding box format is indicated by attribute center_point_box.
Note that this algorithm is agnostic to where the origin is in the coordinate system
and more generally is invariant to orthogonal transformations and translations of
the coordinate system; thus translating or reflections of the coordinate system result
in the same boxes being selected by the algorithm. The selected_indices output
is a set of integers indexing into the input collection of bounding boxes
representing the selected boxes.
The bounding box coordinates corresponding to the selected indices can then be obtained
using the Gather or GatherND operation.

- Version
    This version of the operator has been available since version 11
    of the default ONNX operator set.

- Attributes
    center_point_box : int (default is 0)
        Integer indicate the format of the box data.
        The default is 0.
        0 the box data is supplied as [y1, x1, y2, x2] where (y1, x1) and (y2, x2)
        are the coordinates of any diagonal pair of box corners
        and the coordinates can be provided as normalized
        (i.e., lying in the interval [0, 1]) or absolute.
        Mostly used for TF models. 1 - the box data is supplied as
        [x_center, y_center, width, height]. Mostly used for Pytorch models.

- Inputs (2 - 5)
    boxes : tensor(float)
        An input tensor with shape [num_batches, spatial_dimension, 4].
        The single box data format is indicated by center_point_box.
    scores : tensor(float)
        An input tensor with shape [num_batches, num_classes, spatial_dimension]
    max_output_boxes_per_class (optional) : tensor(int64)
        Integer representing the maximum number of boxes to be selected per batch per class.
        It is a scalar. Default to 0, which means no output.
    iou_threshold (optional) : tensor(float)
        Float representing the threshold for deciding whether boxes overlap
        too much with respect to IOU. It is scalar.
        Value range [0, 1]. Default to 0.
    score_threshold (optional) : tensor(float)
        Float representing the threshold for deciding when to remove boxes based on score.
        It is a scalar.

- Outputs
    selected_indices : tensor(int64)
        selected indices from the boxes tensor.
        [num_selected_indices, 3], the selected index format
        is [batch_index, class_index, box_index].

話が少し逸れましたが、オペレーションをコンポーネント化できる、ということにおいて、PyTorchから実質的に綺麗なオペレーションを生成できない部分を部品化したあとに単純な前後処理をマージしておくことで、いつも定型的にコーディングしている後処理・前処理を事前にストックしておいて再利用できる、というメリットの部分の雰囲気はご理解頂けたのではないでしょうか。ただし注意点は、たとえONNXの仕様上対応しているオペレーションであったとしてもTensorRTなどの他のフレームワーク側では対応していないことがあるという点ぐらいです。

余談ですが、下図はバッチ処理とマルチクラスに対応したNMSをMyriad向けにカスタムしたコンポーネントの一例です。フレームワーク側のバグをONNXのモデル構造でフォローすることができます。MyriadはNMSに対応しているにも関わらず、NMSの出力の型がINT64のままだとAbortするバグがありますので、NMSの直後にダミーの Mul をマージしておくだけでAbortを回避できるトリックがあります。これ以外にも特定のオペレーションの型変換のバグの回避など、挙げ始めるとキリがないためこれぐらいにしておきます。

6. ONNXモデルのチューニング [基礎]

この章では、PyTorchから生成したONNXあるいは直接生成したONNXの構造を機械的に最適化する方法ならびに手動で構造を書き換えて処理フローを変更する方法をご紹介します。

6-1. onnx-simplifierによる最適化

４章でも少しだけ触れた onnx-simplifier の基本動作および最適化の挙動をご紹介します。
コマンドラインインタフェースで使用する場合の基本構文は下記のとおりです。

onnxsim input.onnx output.onnx

たったこれだけを実行することでONNXのモデル構造を大幅に最適化し、場合によってはファイルサイズを大幅に削減してくれます。onnxsim を実行したときに内部で処理してくれる最適化の動きと注意点を下記の流れでご説明します。

1. BatchNormalization の融合
2. オペレーションの定数変換と埋め込み
3. 形状推定
4. 使用するうえでのテクニック
5. Tile と GatherElements のコンビネーションによるモデル肥大化の話

6-1-1. BatchNormalization の融合

BatchNormalization 自体はそんなに複雑な処理をしていません。定型的な処理を担っているだけで、特に推論時にはオペレーションとしてモデルの中に存在している意義があまりありません。ONNXの公式仕様に書かれている BatchNormalization の演算仕様を拝借します。

output = (input_tensor - input_mean) / sqrt(input_var + epsilon) * scale + bias

入力値から平均値を引き算して分散のルートで割り算し、スケールを掛け算してバイアスを足しています。(数式をただ日本語にしただけです）epsilon はゼロ除算エラー防止のための微小な数値です。なお、input_mean と input_var と epsilon と scale と bias はトレーニング終了時点で確定しており全て定数になっている前提です。つまり、input_tensor 以外の数値は全て定数であり、前後のオペレーション(Convolution) などが保持している定数値と事前に足し掛けすることで BatchNormalization オペレーションそのものを無効化することができます。このあたりの考え方は私以外の世界中の方々が昔から解説記事などを沢山書かれていますので、詳しいことは先人の方々が書かれた記事をご覧ください。

ここでお伝えしたいのは、難しい BatchNormalization の前後オペレーションとの融合処理を onnx-simplifier が全て代行して自動で処理してくれる、ということです。下図に例示します。

Before	After

Conv の直後にある BatchNormalization が Conv に融合して無くなりました。しかし、画像右下の BatchNormalization は Conv 周辺には無いため融合も分解もされずそのまま残りました。この点は少し手を加えることで単純な Mul と Sub と Add に分離でき、ときには Mul と Add のみに分離することができますがここでは説明しません。BatchNormalization を分解することで内部的な演算回数を削減することができます。

6-1-2. オペレーションの形状推定と定数変換と埋め込み

事前に演算しておけるオペレーションは onnx-simplifier が定数に置き換えてくれます。前述の図の Unsqueeze や Constant、ConstantOfShape、Tile などです。Vision系のモデルで特に多いのが、Resize 後のサイズ、あるいはリサイズのスケール値を推定する処理部分の Resize オペレーションへの埋め込みです。

Before	After

Resize オペレーションの直前にあった Shape と Slice と Concat の３つのオペレーションから事前にResize形状を計算し、Resize オペレーションの size の Initializer に定数として埋め込まれました。このような動作は Resize に限ったことではなく最適化の過程で事前計算可能な部分は定数に置き換えて埋め込むことでモデル全体の演算量を減らしてくれます。最近流行りの Diffusion モデルなど全体構造が大きなモデルほど効果が大きくなりますが、onnx-simplifier の最適化処理に数時間掛かったり何日も掛かったりすることがあります。

6-1-3. 使用するうえでのテクニック

モデルの構造が極めて大きく、なおかつ形状推定と最適化によって定数化できるオペレショーンの数が非常に多い場合 (例としてモデル全体で数万オペレーションある場合など) は１度の最適化処理で構造を最適化しきれない場合があります。したがって、onnx-simplifier は定型的に３回〜５回実行することをおすすめします。

Hugging Face にコミットされている Stable Diffusion の vae_encoder.onnx を例にとります。合計で３回 onnxsim を実行し、２回目までモデル全体のオペレーションの総数とモデルサイズがシュリンクしていることが分かります。冗長な計算を行っているオペレーションをモデルからほぼ一掃してくれました。

最適化前と最適化後でどれほど差が出ているかをモデル構造を分析することのみに特化したツール ssc4onnx を使用して計測してみます。onnx-simplifier の構造分析部分のロジックを独自に改善し、分析処理を内部的に10倍以上高速化していることと、オペレーションの総数をカウントして表示できるように改良したツールです。onnx-simplifier は巨大な構造を持つモデルを最適化する過程でAbortすることがありますが、このツールはAbortしません。

• コマンドラインインタフェース

  ssc4onnx -h
  
  usage:
      ssc4onnx [-h]
      -if INPUT_ONNX_FILE_PATH
  
  optional arguments:
    -h, --help
          show this help message and exit.
  
    -if INPUT_ONNX_FILE_PATH, --input_onnx_file_path INPUT_ONNX_FILE_PATH
          Input onnx file path.

• Pythonスクリプトから利用する場合

  >>> from ssc4onnx import structure_check
  >>> help(structure_check)
  
  Help on function structure_check in module ssc4onnx.onnx_structure_check:
  
  structure_check(
      input_onnx_file_path: Union[str, NoneType] = '',
      onnx_graph: Union[onnx.onnx_ml_pb2.ModelProto, NoneType] = None
  ) -> Tuple[Dict[str, int], int]
  
      Parameters
      ----------
      input_onnx_file_path: Optional[str]
          Input onnx file path.
          Either input_onnx_file_path or onnx_graph must be specified.
          Default: ''
  
      onnx_graph: Optional[onnx.ModelProto]
          onnx.ModelProto.
          Either input_onnx_file_path or onnx_graph must be specified.
          onnx_graph If specified, ignore input_onnx_file_path and process onnx_graph.
  
      Returns
      -------
      op_num: Dict[str, int]
          Num of every op
      model_size: int
          Model byte size

実に 269個 のオペレーションが最適化の過程で削除されました。

• 最適化前

  

• 最適化後

  

6-1-4. Tile と GatherElements のコンビネーションによるモデル肥大化の話

onnx-simplifier の最適化の仕組みが事前演算によるオペレーションの定数化であることはご紹介しました。ここではその最適化仕様が裏目になってモデルサイズが超肥大化し、Protocol Buffers のファイルサイズ上限の 2GB を超えるほど肥大化してAbortしてしまう事例があることをご紹介します。

モデル肥大化の原因と対処方法を onnx-simplifier の作者と議論し、現状の最新バージョン v0.4.8 では提案内容が取り込まれてモデルサイズの肥大化が抑止される仕様が盛り込まれた状態になっていますが、諸刃の剣の対処であることを理解しつつ、現状ではデフォルト動作で Tile の最適化がスキップされるように修正されています。気になる方はこちらの issue Excessive bloating of ONNX files due to over-efficient conversion of “Tile” to constants (Protocol Buffers .onnx > 2GB) #178 をご覧ください。

ポイントを絞ってご説明すると、Tile オペレーションを定数化する過程で、Int64 形式の定数が 1x16x320x180x320個 (つまり 294,912,000個) 約３億個も生成されてONNXモデルに埋め込まれる爆弾を抱えていました。294,912,000 x 64bit ですので、2.2 GB です。PyTorchの torch.meshgrid を使用した場合に発生する問題です。

画像中の GatherElements は私のワークアラウンドを適用したあとの Int32 へキャストした GatherElements ですので、半分のサイズの 1.1 GB に縮小した状態です。それでもなお、Tile ひとつと GatherElements ひとつの組み合わせ部分だけで 1.1 GB ものモデル肥大化を招きます。onnx-simplifier のデフォルトの動作で Tile の最適化がスキップされるため、Tile を含めた最適化有効・最適化無効を切り替えるオプションの実装が待たれます。

6-2. simple-onnx-processing-tools

基礎編の最後です。次章以降に応用的なモデルチューニングをご紹介するに先立って、onnx-graphsurgeon や onnx-simplifier と組み合わせて使用するONNX加工用ツール群をご紹介します。simple-onnx-processing-tools というツール群です。ONNXモデルの分割、結合、OP削除、サイズ圧縮、属性・定数書き換え、OP生成、OPセット変更、指定入力順への変更、OP追加、RGB→BGR変換、バッチサイズ変更、OP一括リネーム、JSON変換などが可能です。すでに前章までの基礎編でごく一部のツールを活用してONNX構造の基礎理解に役立てましたが、応用編では未登場のツールを使用してあらゆる変換を試行します。

対応しているタスク・ツールは下記のとおりです。応用編で常用します。全て自作のツールですが、ONNX加工に関するほぼ全てのユースケースに対応しています。

No.	Tool Name	Tags	Summary
1	snc4onnx		Simple tool to combine(merge) onnx models. Simple Network Combine Tool for ONNX.
2	sne4onnx		A very simple tool for situations where optimization with onnx-simplifier would exceed the Protocol Buffers upper file size limit of 2GB, or simply to separate onnx files to any size you want. Simple Network Extraction for ONNX.
3	snd4onnx		Simple node deletion tool for onnx. Simple Node Deletion for ONNX.
4	scs4onnx		A very simple tool that compresses the overall size of the ONNX model by aggregating duplicate constant values as much as possible. Simple Constant value Shrink for ONNX.
5	sog4onnx		Simple ONNX operation generator. Simple Operation Generator for ONNX.
6	sam4onnx		A very simple tool to rewrite parameters such as attributes and constants for OPs in ONNX models. Simple Attribute and Constant Modifier for ONNX.
7	soc4onnx		A very simple tool that forces a change in the opset of an ONNX graph. Simple Opset Changer for ONNX.
8	scc4onnx		Very simple NCHW and NHWC conversion tool for ONNX. Change to the specified input order for each and every input OP. Also, change the channel order of RGB and BGR. Simple Channel Converter for ONNX.
9	sna4onnx		Simple node addition tool for onnx. Simple Node Addition for ONNX.
10	sbi4onnx		A very simple script that only initializes the batch size of ONNX. Simple Batchsize Initialization for ONNX.
11	sor4onnx		Simple OP Renamer for ONNX.
12	soa4onnx		Simple model Output OP Additional tools for ONNX.
13	sod4onnx		Simple model Output OP Deletion tools for ONNX.
14	ssi4onnx		Simple Shape Inference tool for ONNX.
15	sit4onnx		Tools for simple inference testing using TensorRT, CUDA and OpenVINO CPU/GPU and CPU providers. Simple Inference Test for ONNX.
16	onnx2json		Exports the ONNX file to a JSON file.
17	json2onnx		Converts a JSON file to an ONNX file.
18	sed4onnx		Simple ONNX constant encoder/decoder. Since the constant values in the JSON files generated by onnx2json are Base64-encoded values, ASCII <-> Base64 conversion is required when rewriting JSON constant values.
19	ssc4onnx		Checker with simple ONNX model structure. Simple Structure Checker for ONNX. Analyzes and displays the structure of huge size models that cannot be displayed by Netron.
20	sio4onnx		Simple tool to change the INPUT and OUTPUT shape of ONNX.
21	svs4onnx		A very simple tool to swap connections between output and input variables in an ONNX graph. Simple Variable Switch for ONNX.
22	onnx2tf		Self-Created Tools to convert ONNX files (NCHW) to TensorFlow format (NHWC). The purpose of this tool is to solve the massive Transpose extrapolation problem in onnx-tensorflow (onnx-tf).
23	sng4onnx		A simple tool that automatically generates and assigns an OP name to each OP in an old format ONNX file.
24	sde4onnx		Simple doc_string eraser for ONNX.
25	components_of_onnx	[WIP]	ONNX parts yard. The various operations described in Operator Schemas are converted in advance into OP stand-alone ONNX files.

X. おわりに

ONNXのモデルチューニングテクニック (基礎編) は以上です。次回は 応用編１ として実践的なONNXモデル加工のテクニックをご紹介したいと思います。

7. モデルのチューニング [応用]

AI Lab Creative Researchグループでリサーチエンジニアをしている大田（@ciela）です。主にAI LabリサーチサイエンティストやAI事業本部プロダクト開発チームに対する機械学習エンジニアリングサポートを行っています。

昨年末、弊社よりFEED自動承認AIというサービスがリリースされました。このサービスはGoogle Merchant Center（GMC）にて不承認となった商品画像から不承認理由に該当する領域を除去することで、承認率を向上することのできるサービスです。内部では機械学習モデルが利用されており、不承認画像それぞれに特化した画像編集が施される仕組みとなっています。詳しくは下記プレスリリースをご参照ください。

AIによってGoogleショッピング広告における承認率の大幅改善を図るサービス「FEED自動承認AI」を開発・提供開始

FEED自動承認AIのモデル開発にはAI Labからも下田と私大田が関わっているのですが、主に下田がモデルの開発と学習を行い、私がシステム化してプロダクトチームへ提供する、といった役割分担でプロダクトサポートを進めています。下田はリサーチサイエンティストとしてモデル精度向上のための試行錯誤を継続していく必要があり、並行してシステムへモデルコードをデプロイしていくことがなかなか難しい立場です。しかし、サービスとして実用化するためにはモデルの性能がある程度見極められたタイミングでデプロイしていかなければなりません。そこで私がシステム化に関するエンジニアリングサポートに入ることとなりました。

本記事では、FEED自動承認AIリリースまでに私が行ってきた実験用モデルをシステム導入する上での取り組みについて記したいと思います。ほぼPyTorchとAWS SageMakerに特化した内容ではあるものの、研究で組み立てている機械学習モデルを実用化していく上での経験談としてお役に立てれば幸いです。

モデルリファクタリング

今回のFEED自動承認AIモデルの入力・出力は下記のとおりです:

• 入力: GMCにて不承認となった商品画像
• 出力: 不承認理由に該当する領域が編集・除去された画像

モデル内部ではさらにいくつか特定の処理に特化したサブモデル・画像処理で構成されており、PyTorch、OpenCV、Pillowを利用して実装されています。重みファイルは3つあり合計850MBほどの容量です。実験用にほぼ一直線になっていたコードをモジュール・クラスに分割して構造化し、推論に特化した形に仕立てあげていったのですが、その際大事にしていたポイントをまとめてみます。

不要な計算・重みの削除

モデルコードを編集していくにあたり、まずは無駄なレイヤーや計算が含まれていないかを確認するところからはじめました。PyTorchの場合、

nn.Module

nn.Module.load_state_dict

strict=False

モデルコードと重みファイル内の重みデータに過不足がないようにするために、モデルクラスから不要なレイヤーを削除した後、一度

strict=False

torch.save

strict=False

CPU/GPU間データ移動の除去

GPU上のモデルでデータを処理する際、理想的には一度GPUに送ったデータは出力時までずっとGPU上で計算させたいものです。今回のモデルにおいては、サブモデル同士の出力と入力の間にOpenCVとPillowを用いた画像処理が挟まれており、このためにCPU/GPU間でのデータ移動が発生していました。そのため、この処理をGPU上で行うように置き換えデータ移動が発生しないようにできないかを考えました。

嬉しいことに最近のtorchvisionではOpenCVやPillowとほぼ似たようなパラメータを扱えるAPIで

torch.Tensor

システムへのデプロイ

FEED自動承認AIのシステムはAWS上に構築されています。そのため、今回のモデルはAWSの機械学習統合プラットフォームであるSageMakerにデプロイしていくことになりました。SageMakerは以前から何度か研究用途で利用してはいたものの、プロダクトシステムへと提供される推論エンドポイント構築のために利用するのは以前参加したワークショップ以来でした。

学習済みモデルの推論エンドポイント化

モデルの開発・学習からSageMaker上で行っている場合は推論エンドポイントへのデプロイが非常にスムーズなことをワークショップでも体験させてもらっていたのですが、今回は学習済みモデルが既に手元にある状態です。そのため、下記のSageMakerドキュメントページに記載されている外部モデルのデプロイ方式を採用しようと考えました。

I have a PyTorch model that I trained outside of SageMaker, and I want to deploy it to a SageMaker endpoint
Use PyTorch with Amazon SageMaker

しかしながら、2022年2月現在、このデプロイ方式に記されているリンク先でのアンカリングが効いておらず、ページ内の該当項目にナビゲートされないという状況に陥ってしまいます…。おそらくリンク先のSphinxドキュメントのみ更新されたままなのでしょう。手探りで色々調べたり普段からお世話になっているAWS Solution Architectの方々に尋ねたりしつつ、下記のImmersion Dayのリポジトリに含まれているsickit-learnモデル向けのハンズオンノートブックに従って作業することで、なんとか推論エンドポイントとしてデプロイすることができました。

Lab: Bring your own custom container with Amazon SageMaker

もし今回のように既に学習済みモデルがある状態でのSageMaker推論エンドポイントへのデプロイを考えている方は参考にしていただけたらと思います。この作業において重要だと感じたポイントを下記に記しておきます。

• 重みファイルは

  model.tar.gz

   という名前でアーカイブしてS3にアップロードしておく

• serve

  というファイル名のスクリプトを実行可能状態でDockerイメージ内に含めパスを通す（今回は推論なので

  serve

   のみだが学習の場合は

  train

   も準備する）
• ヘルスチェック用の

  /ping

   と実際の推論を行う

  /invocations

   の2種類のAPIを準備したWebサーバを記述し

  serve

   スクリプトから起動する

ローカルでの動作確認

推論エンドポイントとしてのモデルの動作確認をWebリクエスト経由で行う場合、最終的にはSageMakerにデプロイして確認しなければなりません。上記Immersion Dayのノートブック通りに作業した場合、FlaskとGunicornを使ったWebサーバが推論エンドポイントとなりますが、それらの設定ミスなどの些細なことで発生したエラーによってデプロイにかかった時間を無駄にするのはもったいありません。

SageMakerではローカルモードというローカルマシン上にSageMaker環境を構築できる機能も提供しています。SageMakerにデプロイする前には必ずこのローカルモードで推論エンドポイントの動作確認を事前に行うことで、デプロイ後の手戻りをなるべく少なくするように心がけていました。

ノートブックインスタンスのトレーニングに Amazon SageMaker のローカルモードを使用する
オンプレミス環境から Amazon SageMaker を利用する

その他

基礎的ですがシステムへの導入にあたって他に考慮していた項目を挙げておきます:

• モデル推論時には

  nn.Module.eval

   で推論モードにし、

  torch.no_grad

   でガードして勾配計算をカットする
• SageMakerインスタンスタイプに合わせてGunicornパラメータ（ワーカー数、スレッド数など）を調節する
• SageMakerバッチ推論における同時並列リクエスト数を調節する

特にGunicornパラメータ周りについては、パフォーマンス確認の際にいちいちコードから変更する必要がないように、SageMaker推論エンドポイント起動時に渡すことのできる環境変数値を介して調節が出来るようにスクリプトを記述しました。

おわりに

今回の記事ではFEED自動承認AIのリリースに向けて、私が実験用コードを推論エンドポイントとして実用化するまでの取り組みについてまとめてみました。下田とうまく分業してAI Labから事業への貢献事例をつくれたのではないかと思っています。

ちなみにJIT化による高速化などモデル自体のパフォーマンス向上や、PyTorch推論アプリケーションサーバであるTorchServeについての調査など、今後FEED自動承認AIモデルの推論エンドポイントに関して検証したいタスクはたくさんあります。もしAI Labにてプロダクトサポートや研究のためのエンジニアリングに興味のある方がいらっしゃれば、どうぞお気軽にご連絡いただければと思います。

深層学習データセット取扱の課題

昨今の深層学習に用いられるデータセット内のデータ数は一般的に大規模です。実際に学習する上ではデータセットをシャッフルしつつ繰り返しロードし、場合によっては適切な前処理やData Augmentationだってしなくてはなりません。それらの処理を並列化することで、学習にかかる時間を少しでも減らそうと涙ぐましい努力をされている方も多いことでしょう。その一方で、時には数万件にも及ぶファイルに対してランダムアクセスしつつ何度もロードするアクセスパターンは入出力効率を低下させてしまうためストレージデバイスにとって優しくありません。この矛盾を解決するために、格納プロトコルによって規格化された独自アーカイブファイルにデータセットを変換し、それを一定サイズに分割（シャーディング）した形で保管するアプローチが採用されたりします。ある程度まとまった容量のファイルをシーケンシャルにロードすることで、ストレージの入出力効率を高めることが狙いです。

深層学習フレームワークの代表格であるTensorFlowにおいてはプロトコルバッファ形式でデータをシリアライズ・アーカイブするTFRecordという仕組みが提供されています。例えばGoogle Cloud Platform (GCP)のCloud Storageに200〜300MB程度の容量ごとにシャーディングされた学習データのTFRecordファイルを配置しておけば、あとはそのパスパターンを指定するだけで学習データをシーケンシャルにロードすることができ、ファイル単体ごとにランダムにロードする場合と比べて学習速度の向上が期待できます。学習データの前処理パイプライン構築に最適なApache Beam（GCP Dataflow）からも簡単に書き出せるため、GCP上でTensorFlowを利用している方にとっては既に当たり前の仕組みかもしれません。一方で、TensorFlowと肩を並べる深層学習フレームワークであるPyTorchではどういったアプローチが考えられているのでしょうか。昨年8月、PyTorch公式ブログにEfficient PyTorch I/O library for Large Datasets, Many Files, Many GPUsという記事が投稿されており、記事内ではWebDatasetというTensorFlowでのTFRecordに相当する仕組みが紹介されています。本記事ではこのWebDatasetについて簡単にまとめ、実際のデータセットでの利用例をお見せしようと思います。

WebDataset

WebDatasetは国際会議BigData 2019にて採択されたHigh Performance I/O For Large Scale Deep Learningという論文で定義されている大規模な深層学習データセットを取り扱うための機構およびその実装です。PyTorchのIterableDatasetクラスを継承実装しており、データベース・リモートサーバ・リアルタイムログからストリーミングでデータを取得するような場合など、データ全体に対するランダムアクセスが困難・不可能でバッチサイズがロード済みデータに依存してしまうような環境下でのデータセットを取り扱うことに適した仕組みになっています。TensorFlowのTFRecord同様、WebサーバやクラウドストレージにWebDatasetとしてデータをシャーディングして配置しておくことでその状態のまま学習ルーチンへ取り込むことが可能となります。2021年3月現在はPyTorchとは別リポジトリで開発が進められているものの、既にRFCとしてプロポーザルが上がっており将来的にはPyTorch内部に取り入れられるかもしれません。

WebDatasetの大きな特徴として、実体がPOSIX tar形式でアーカイブされていることが挙げられます。*nixなコンピュータを利用する方にとっては馴染み深すぎる形式ですね。データセット内の各データとラベルを同一のbasenameを持ったファイルとして対応付けた状態でアーカイブしており、tarコマンドで中身を確認したり展開したりといったことも容易に行なえます。読み込み用の特別な仕組みを必要としないため、ポータブルなデータセットアーカイブとしての役割も果たしていると言えるでしょう。下記は手元で作成したWebDatasetの中身を閲覧してみた例です。

$ tar tvf bam-wds-0001.tar | head -n 4
-r--r--r-- bigdata/bigdata 14     2021-03-30 12:02 content_flower/4.attr # ラベルファイル
-r--r--r-- bigdata/bigdata 442991 2021-03-30 12:02 content_flower/4.jpg  # データファイル
-r--r--r-- bigdata/bigdata 14     2021-03-30 12:02 content_flower/5.attr # ラベルファイル
-r--r--r-- bigdata/bigdata 176634 2021-03-30 12:02 content_flower/5.jpg  # データファイル

WebDatasetを読み込んだあとの、シャッフル・torch.Tensor化・バッチ化などの学習中前処理はTensorFlow tf.data.Datasetと似たような書き方で定義できます。具体的にはメソッドチェーンを利用して処理を繋げていくfluent interfaceとなっており、下記のような形で記述できます。

# 100件のバッファでデータをシャッフル、RGB3チャンネルのtorch.Tensorとして実体化、バッチサイズ4でバッチ化
wds.WebDataset("bam-wds-0001.tar").shuffle(100).decode("torchrgb").batched(4)

この他にも、メディアファイルの拡張子に基づく自動デコード・エンコード、取得したリソース（シャード、レコード共に）のローカルキャッシュ、などアーカイブされたデータセットを取り扱う上で嬉しい工夫が凝らされてます。詳しくはGitHubリポジトリのJupyter Notebook集によくまとまっていますのでこちらをご参照ください。

既存のデータセットをWebDatasetへ変換

実際に既存のデータセットをWebDatasetへ変換してみたいと思います。例としてBehance Artistic Media（BAM）データセットを利用してみましょう。BAMデータセットの実体はSQLiteファイルとなっており、DBレコードのsrcカラムに指定されたBehance CDNのURLを参照することで画像をロードできますが、学習時にいちいちインターネットアクセスをしていては学習ルーチンの効率が悪すぎます。さらに既にアクセスが不可能となってしまっている（404）画像に対する例外処理が必要であったりと、そのままでは扱いが若干ややこしいデータセットです。画像を予めダウンロードしてローカルストレージに保存しておき、各DBレコードのURLをファイルパスへ変換しつつランダムアクセスするのがPyTorchでの素直な学習ワークロードになるかと思われますが、ここではそれをWebDatasetに置き換えてみましょう。

今回は下記のようなSQLを発行して得られる、BAMデータセットDBのcrowd_labelsテーブルによる属性ラベルを画像と対応付けたレコードをデータセットとして考え、これをWebDataset化したいと思います。BAMデータセットの主要なDBスキーマについては以前の記事で簡単に紹介しておりますのでよろしければご覧ください。

クエリ

SELECT
  *
FROM
  (
    SELECT
      ROW_NUMBER() OVER(PARTITION BY attribute) AS attrseq, -- ラベルごとの通し番号
      src, -- 画像URL
      attribute -- ラベル
    FROM
      modules
      INNER JOIN crowd_labels
      ON  modules.mid = crowd_labels.mid
    WHERE
      label = 'positive'
  )
WHERE
  attrseq <= 10

レコードサンプル

1|https://mir-s3-cdn-cf.behance.net/project_modules/disp/e0533d11631403.560fad7cb723b.jpg|content_bicycle
2|https://mir-s3-cdn-cf.behance.net/project_modules/disp/88e45042270399.5607181dccf11.jpg|content_bicycle
3|https://mir-s3-cdn-cf.behance.net/project_modules/disp/72053821069043.562fb004d1c67.jpg|content_bicycle
4|https://mir-s3-cdn-cf.behance.net/project_modules/disp/37cbc824014175.5632ca1c9f9bb.jpg|content_bicycle
5|https://mir-s3-cdn-cf.behance.net/project_modules/disp/5cd51b922413.5629a98a73644.jpg|content_bicycle
6|https://mir-s3-cdn-cf.behance.net/project_modules/disp/e84475922413.5629a93f2986e.jpg|content_bicycle
7|https://mir-s3-cdn-cf.behance.net/project_modules/disp/7456a329931731.560b01dbde761.jpg|content_bicycle
8|https://mir-s3-cdn-cf.behance.net/project_modules/disp/a6c90115285239.5628f21d09ca1.jpg|content_bicycle
9|https://mir-s3-cdn-cf.behance.net/project_modules/disp/80cdd024014175.5632ca1caebcf.jpg|content_bicycle
10|https://mir-s3-cdn-cf.behance.net/project_modules/disp/8aae9918574135.562cbbaa0c1d3.jpg|content_bicycle
1|https://mir-s3-cdn-cf.behance.net/project_modules/disp/02a0e610386217.560e418f245eb.JPG|content_bird
2|https://mir-s3-cdn-cf.behance.net/project_modules/disp/01d54d10386217.560e4167dc493.JPG|content_bird
3|https://mir-s3-cdn-cf.behance.net/project_modules/disp/eb2a8d10204191.560e10d04ed41.jpg|content_bird
4|https://mir-s3-cdn-cf.behance.net/project_modules/disp/c4b5d220463133.562ebb165443f.jpg|content_bird
5|https://mir-s3-cdn-cf.behance.net/project_modules/disp/6b1b1120463133.562ebb165faa8.jpg|content_bird
（以下略）

今回は少量データで試すために各属性ごとのデータは10件ずつに絞っています。さらにこのデータセットをWebDatasetに変換するために下記のようなコードを実行してみましょう。

db_file = "20170509-bam-1m-18UThu3ICM.sqlite"
wds_pattern = "bam-wds-%04d.tar.gz"
with sqlite3.connect(db_file) as con,\
    wds.ShardWriter(wds_pattern, maxsize=10_000_000, compress=True) as writer:
    cur = con.cursor()
    cur.execute(QUERY)  # 上記クエリ
    for row in cur:
        url = row[1]
        try:
            _, ext = os.path.splitext(url)
            with urlopen(url) as res:
                body = res.read()
            sample = {
                "__key__": f"{row[2]}/{row[0]}",  # [attr]/[attrseq]
                ext[1:].lower(): body,            # png or jpg
                "attr": row[2],
            }
            writer.write(sample)
            log.debug(f"Archived as WDS {row}")
        except HTTPError as e:
            log.warning(f"{url}: {e.code}")
        except URLError as e:
            log.warning(f"{url}: {e.reason}")

URLにアクセスして取得できたバイト列をリソースの拡張子に合わせてWebDatasetに一件一件格納していきます。さらにシャーディングを試すために圧縮前の最大ファイルサイズはあえて10MBに制限しています（デフォルト3GB）。実行した結果、手元の環境に4つのtarballファイルが生成されました。データセットのWebDataset化自体はこれで完了です。簡単ですね。

$ ls bam-wds-*.tar.gz
bam-wds-0000.tar.gz bam-wds-0001.tar.gz bam-wds-0002.tar.gz bam-wds-0003.tar.gz

あとは読み込みです。上記のローカルファイルをファイルパスとして直接読み込んでも良いのですが、せっかくなのでこれらが外部サーバに配置された状態を想定し、HTTP経由でウェブリソースとして取得するシミュレーションを行ってみましょう。Python3のhttp.serverモジュールを利用して簡易ウェブサーバを立上げWebDatasetをローカル配信してみます。

$ python3 -m http.server

この状態で下記のコードを実行してみましょう。WebDatasetに格納されているデータセットから、画像はRGB3チャンネルでサイズ224×224のtorch.Tensor、ラベルはstrとしてそれぞれデコード・前処理を行って読み込むコードになります。

wds_urls = "http://localhost:8000/bam-wds-{0000..0003}.tar.gz"
dataset = (
    wds.WebDataset(wds_urls, shardshuffle=True)
    .shuffle(100)
    .decode("torchrgb")
    .to_tuple("jpg;png", "attr")
    .map_tuple(transforms.Resize((224, 224)), lambda x: x.decode())
    .batched(4)
)
loader = data.DataLoader(dataset, num_workers=4, batch_size=None)
for tensors, attrs in loader:
    log.info(f"Tensor: {tensors.size()} Labels: {attrs}")

実行結果は下記のようになりました。無事WebDatasetを読み込むことに成功し、画像とその属性ラベルがバッチ化された状態で取得できたことが確認できました。

Tensor: torch.Size([4, 3, 224, 224]) Labels: ['content_cars', 'content_bicycle', 'content_bicycle', 'content_building']
Tensor: torch.Size([4, 3, 224, 224]) Labels: ['content_tree', 'emotion_peaceful', 'emotion_peaceful', 'content_flower']
Tensor: torch.Size([4, 3, 224, 224]) Labels: ['emotion_scary', 'emotion_scary', 'media_comic', 'media_3d_graphics']
Tensor: torch.Size([3, 3, 224, 224]) Labels: ['media_watercolor', 'media_watercolor', 'media_watercolor']
Tensor: torch.Size([4, 3, 224, 224]) Labels: ['content_building', 'content_bird', 'content_bird', 'content_bicycle']
（以下略）

おわりに

本記事では深層学習データセットを取り巻く課題とWebDatasetによるアプローチを簡単に解説しました。PyTorchで大規模なメディアデータセットを取り扱う上での指針となるのではないでしょうか。

最近ではTensorFlow Datasets（TFDS）というライブラリも出てきており、既存のデータセットをより簡単にTFRecordとして読み書きする仕組みが整っていたりします。こちらは名前の通りTensorFlow向けの仕組みですが、同様の思想のもとPyTorch向けにWebDatasetを書き出すような何らかのシステムを整えることは可能なはずです。独自データセットをTraining-readyな状態で利用したいのにTensorFlow向けにしか提供されていない…というような危険な宗教戦争の火種を潰せますね。

AI Labではこのような研究のためのエンジニアリングをより加速させていきたいと思います。今後は採用も考えていければと思うのでこういった活動に興味のある方はどうぞご連絡ください。

AI Labの大田（@ciela）です。近頃の機械学習モデルにおいて、敵対的な学習を行うことで有名なのは生成器ネットワークと識別器ネットワークが互いに競い合うように学習するGenerative Adversarial Networks (GAN) かと思いますが、転移学習の一種であるドメイン適応においても敵対的な学習を行うモデルが提案されています。ドメイン適応の目的は、学習データ（ソースドメイン）とテストデータ（ターゲットドメイン）におけるデータ分布が異なるような状況においてもうまく機能するようなモデルを学習することです。ターゲットドメインのデータ分布をソースドメインのデータ分布に近づけるように学習するのですが、これを敵対的なやり方で行うというわけです。ドメイン適応自体についてはAI Labの安井もブログ記事を書いていますのでこちらも併せてお読みいただければ幸いです。

Ganinらは教師なしドメイン適応モデルであるDomain-Adversarial Neural Network (DANN)を提案しており、異なる画像表現を持つ画像データセットでの有効性を示しています。

Domain-Adversarial Training of Neural Networks [Ganin, JMLR2016]

論文よりDANNのアーキテクチャ図

上図を参照してみると特徴抽出CNNを経たあとクラス分類とドメイン分類の2つのサブネットワークに分岐していますが、ドメイン分類のサブネットワークにおいてはGradient Reversal Layer (GRL)という層を通していることが分かります。この層は誤差逆伝搬時にのみ前段の層へ伝わる勾配の符号を反転させるのですが、これによってクラス分類能力を高めつつもドメイン分類能力は落とすという敵対的な学習が行われます。その結果、前段の特徴抽出層においてドメイン不変性のある特徴獲得能力が向上するという仕組みになっています。

このGRLはGANの学習にも利用できます。生成器と識別器を交互に学習させる実装が主流かと思いますが、下記の記事においてはGRLを利用して生成器へ伝搬される勾配を反転するだけで目的関数がシンプルになり実装が簡単になるという利点が紹介されています。

Gradient reversal layerを使うとGANの実装が簡単になる – Qiita

さて、今回は私もドメイン適応モデルを構築する機会があり、ベースモデルとしてDANNを実装する上でGRLについても理解・利用する必要性が出てきました。普段使いしているDNNフレームワークはPyTorchなのですが、2020年1月現在の最新版である1.4.0にはこのGRLの実装は含まれていないようです。これまで活性化関数などを独自実装した経験がなかったこともあり、本記事ではPyTorchでGRLの実装を行い実際に簡単なモデルに組み込んで挙動を確認してみたいと思います。

PyTorchでのGRL実装

PyTorch において、誤差逆伝播時に勾配情報を制御するためには

torch.autograd.Function

PyTorch: Defining New autograd Functions — PyTorch Tutorials 1.4.0 documentation

上記ページを参考にして定義してみた場合、下記のようになります。

class GradientReversalFunction(torch.autograd.Function):
    
    @staticmethod
    def forward(ctx, input_forward: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
        ctx.save_for_backward(scale)
        return input_forward

    @staticmethod
    def backward(ctx, grad_backward: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        scale, = ctx.saved_tensors
        return scale * -grad_backward, None

順伝播時の挙動を

forward

backward

forward

forward

backward

None

GRLの関数実装自体は以上で完了です。簡単ですね。PyTorchでの独自関数実装の入門に最適な題材かと思います。あとはお作法的な問題ではあるのですが、次に行うネットワーク実装の見通しを良くするために

torch.nn.Module

class GradientReversal(nn.Module):

    def __init__(self, scale: float):
        super(GradientReversal, self).__init__()
        self.scale = torch.tensor(scale)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return GradientReversalFunction.apply(x, self.scale)

これでPyTorchでのニューラルネットワークからGRLを利用する準備が整いました。次は実際にGRLの有無での挙動の違いを確認してみましょう。

GRLの挙動確認

ここではDANNモデルのネットワークアーキテクチャを抽象化した下記のようなシンプルなニューラルネットワークを考えます。

入力として任意のベクトルを受け取り共通層を経たあとTrunk層とBranch層へと分岐するのですが、Branch層のみ直前にGRLを配置することでBranch側から逆伝搬された勾配のみ符号反転を行うアーキテクチャです。また、GRLの有無に伴う逆伝播の確認を簡単にするため、Trunk層とBranch層における出力値への影響が等しくなるように両者間での重みを共有するSiamese構造にしています。GRL以外の隠れ層は入力・出力テンソルの次元が全て1次元でバイアス項なしの全結合層とし、GRLのスケーリング係数は1.0を設定しておきます。GRLが正しく機能していれば、Trunk側から逆伝搬されてくる勾配が相殺され、共通層に届く勾配値が0となるはずです。

このアーキテクチャの実装コードは下記になります。

class BranchReversalSiameseNet(nn.Module):

    def __init__(self, use_grl: bool = True):
        super(BranchReversalSiameseNet, self).__init__()
        self.use_grl = use_grl
        self.shared = nn.Linear(1, 1, bias=False)
        self.trunk_branch = nn.Linear(1, 1, bias=False)
        self.grl = GradientReversal(scale=1.0)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        y_shared = self.shared(x)
        y_trunk = self.trunk_branch(y_shared)
        if self.use_grl:
            y_shared = self.grl(y_shared)
        y_branch = self.trunk_branch(y_shared)
        return y_trunk, y_branch

それではGRLの有無による誤差逆伝播を実際に確かめて行きたいと思います。上記クラスをインスタンス化し、ランダム値を入力してTrunk層とBranch層からの出力を合算した値から誤差逆伝播を行い、共通層へ届いた勾配情報を表示するようなスクリプトを記述してみました。最初にネットワーク初期化時の各層の重みも確認のために表示しています。

print('===== BranchReversalSiameseNet Parameters ======')
torch.manual_seed(0)
model = BranchReversalSiameseNet()
print(f'Shared Weight: {list(model.shared.parameters())[0].data}')
print(f'Trunk/Branch Weight: {list(model.trunk_branch.parameters())[0].data}')
print()

print('===== BranchReversalSiameseNet w/o GRL =========')
torch.manual_seed(0)
model = BranchReversalSiameseNet(use_grl=False)
x = torch.randn(1)
print(f'Input: x={x.data}')
y_trunk, y_branch = model(x)
y_out = y_trunk + y_branch
print(f'Output: y_trunk={y_trunk.data}, y_branch={y_branch.data}, y_out={y_out.data}')
y_out.backward()
print(f'Shared Gradient: {list(model.shared.parameters())[0].grad}')
print()

print('===== BranchReversalSiameseNet w/ GRL =========')
torch.manual_seed(0)
model = BranchReversalSiameseNet()
x = torch.randn(1)
print(f'Input: x={x.data}')
y_trunk, y_branch = model(x)
y_out = y_trunk + y_branch
print(f'Output: y_trunk={y_trunk.data}, y_branch={y_branch.data}, y_out={y_out.data}')
y_out.backward()
print(f'Shared Gradient: {list(model.shared.parameters())[0].grad}')
print()

自分の環境において、上記スクリプトの実行結果は下記になりました。いずれも順伝播時の出力自体は等しいものの、GRLありの場合は誤差逆伝播後に共通層へ伝搬されてきた値は0になっており、GRLを利用していない場合とで違いが現れたことが確認できました。

===== BranchReversalSiameseNet Parameters ======
Shared Weight: tensor([[-0.0075]])
Trunk/Branch Weight: tensor([[0.5364]])

===== BranchReversalSiameseNet w/o GRL =========
Input: x=tensor([-0.2905])
Output: y_trunk=tensor([0.0012]), y_branch=tensor([0.0012]), y_out=tensor([0.0023])
Shared Gradient: tensor([[-0.3117]])

===== BranchReversalSiameseNet w/ GRL ==========
Input: x=tensor([-0.2905])
Output: y_trunk=tensor([0.0012]), y_branch=tensor([0.0012]), y_out=tensor([0.0023])
Shared Gradient: tensor([[0.]])

共通層に逆伝播されてくる勾配の値を手計算でも求めてみたいと思います。出力\(y_{\mathrm{out}}\)に対する共通層の重み\(w_{\mathrm{shared}}\)での偏微分\(\frac{\partial y_{\mathrm{out}}}{\partial w_{\mathrm{shared}}}\)を連鎖律を使って展開し、実行時に表示されていた入力値と各層の重みを代入してみましょう。

GRLなし

$$
\begin{eqnarray}
\frac{\partial y_{\mathrm{out}}}{\partial w_{\mathrm{shared}}}
&=& \frac{\partial y_{\mathrm{out}}}{\partial y_{\mathrm{trunk}}} \cdot \frac{\partial y_{\mathrm{trunk}}}{\partial w_{\mathrm{shared}}}
+ \frac{\partial y_{\mathrm{out}}}{\partial y_{\mathrm{branch}}} \cdot \frac{\partial y_{\mathrm{branch}}}{\partial w_{\mathrm{shared}}}\nonumber \\
&=& \frac{\partial y_{\mathrm{out}}}{\partial y_{\mathrm{trunk}}} \cdot \frac{\partial y_{\mathrm{trunk}}}{\partial y_{\mathrm{shared}}} \cdot \frac{\partial y_{\mathrm{shared}}}{\partial w_{\mathrm{shared}}}
+ \frac{\partial y_{\mathrm{out}}}{\partial y_{\mathrm{branch}}} \cdot \frac{\partial y_{\mathrm{branch}}}{\partial y_{\mathrm{shared}}} \cdot \frac{\partial y_{\mathrm{shared}}}{\partial w_{\mathrm{shared}}}\nonumber \\
&=& 1 \cdot w_{\mathrm{trunk}} \cdot x + 1 \cdot w_{\mathrm{branch}} \cdot x\nonumber \\
&=& 0.5364 \cdot (-0.2905) + 0.5364 \cdot (-0.2905)\nonumber \\
&=& -0.3116
\end{eqnarray}
$$

GRLあり

$$
\begin{eqnarray}
  \frac{\partial y_{\mathrm{out}}}{\partial w_{\mathrm{shared}}}
&=& \frac{\partial y_{\mathrm{out}}}{\partial y_{\mathrm{trunk}}} \cdot \frac{\partial y_{\mathrm{trunk}}}{\partial w_{\mathrm{shared}}}
+ \frac{\partial y_{\mathrm{out}}}{\partial y_{\mathrm{branch}}} \cdot \frac{\partial y_{\mathrm{branch}}}{\partial w_{\mathrm{shared}}}\nonumber \\
&=& \frac{\partial y_{\mathrm{out}}}{\partial y_{\mathrm{trunk}}} \cdot \frac{\partial y_{\mathrm{trunk}}}{\partial y_{\mathrm{shared}}} \cdot \frac{\partial y_{\mathrm{shared}}}{\partial w_{\mathrm{shared}}}
+ \frac{\partial y_{\mathrm{out}}}{\partial y_{\mathrm{branch}}} \cdot \frac{\partial y_{\mathrm{branch}}}{\partial y_{\mathrm{GRL}}} \cdot \frac{\partial y_{\mathrm{GRL}}}{\partial y_{\mathrm{shared}}} \cdot \frac{\partial y_{\mathrm{shared}}}{\partial w_{\mathrm{shared}}}\nonumber \\
&=& 1 \cdot w_{\mathrm{trunk}} \cdot x + 1 \cdot w_{\mathrm{branch}} \cdot (-\lambda) \cdot x\nonumber \\
&=& 0.5364 \cdot (-0.2905) + 0.5364 \cdot (-1) \cdot (-0.2905)\nonumber \\
&=& 0
\end{eqnarray}
$$

ちゃんとコード実行時に表示された勾配値が得られました。GRLありの途中式2項目、Branch側の逆伝播に含まれている符号反転スケーリング係数\(-\lambda\)がポイントですね。

おわりに

本記事ではPyTorchでGRLの実装を行い、実際にモデルに組み込んで誤差逆伝播時の挙動を確認しました。これで無事DANNの実装へ進むことができそうで一安心です。今回のコードの全体については下記のGistにまとめてありますので、よろしければこちらもご参照ください。

https://gist.github.com/ciela/2656510f15d2a1e7c5cef206f8328cb2

ちなみに今回の件はサイバーエージェントAI事業本部にて本業の20%の時間を自分の好きな研究活動に費やすことができるゼミ制度での取り組みに端を発するものです。半期ごとに成果目標を掲げて活動計画を提出することで運営から予算を取得し、本業ではなかなか触れることの出来ないデータやモデル検証に挑戦する時間を業務中に得ることが出来る貴重な制度となっています。