【モデルデプロイ】入門チュートリアル(3)：PyTorchからONNXまで詳しく解説

モデル展開入門チュートリアル (3): PyTorch から ONNX への詳細な説明 - Zhihu (zhihu.com)

前の2 つのチュートリアルでは、全員が最初のモデルを正常にデプロイできるように導き、モデルのデプロイメントで発生する可能性のあるいくつかの問題を解決しました。今日からONNX関連の知識を浅いところから深いところまで紹介していきます。ONNX は現在、モデル展開のための最も重要な中間表現の 1 つです。ONNX の技術的な詳細を学習すると、モデルの展開に関する多くの問題を回避できます。
PyTorch モデルを ONNX モデルに変換する場合、多くの場合、1 つの文を簡単に呼び出すだけで済みますtorch.onnx.export。この関数のインターフェースはシンプルに見えますが、その使用には多くの「隠れたルール」があります。このチュートリアルでは、PyTorchモデルをONNXモデルに変換する原理と注意点を詳しく紹介します。さらに、PyTorch と ONNX の間のオペレーターの対応関係も紹介し、PyTorch モデルの変換中に発生する可能性のあるオペレーターサポートの問題に対処する方法を説明します。

プレビュー: 次の記事では、PyTorch でより多くの ONNX オペレーターをサポートする方法を引き続き紹介し、誰もが PyTorch から ONNX への展開ルートを完全に通過できるようにし、ONNX 自体の知識を紹介し、ONNX を変更およびデバッグします。モデルに対する共通のアプローチにより、ONNX に関連する展開上の問題のほとんどを自分で解決できます。乞うご期待~

`torch.onnx.export` 壊す

このセクションでは、PyTorch から ONNX への変換機能を詳しく紹介します torch.onnx.export。このモデル変換インターフェイスをより柔軟に使用し、その実装原理を理解して、この関数のエラーに適切に対処できることを願っています (モデルのデプロイメントの互換性のため、この関数は複雑なモデルをデプロイするときにエラーを報告することがよくあります)。

計算グラフのエクスポート方法

TorchScript は、 PyTorch モデルをシリアル化および最適化するための形式であり、最適化中にtorch.nn.Moduleモデルは TorchScript モデルに変換されます torch.jit.ScriptModule。現在、TorchScript は中間表現としてもよく使用されます。TorchScript の詳細な紹介は他の記事 ( TorchScript の解釈 (1): 初めて TorchScript を知る - Zhihu ) で説明しています。ここでの TorchScript の紹介は、PyTorch モデルを ONNX に変換する原理を説明するためにのみ使用されます。
torch.onnx.exportで必要なモデルは実際にはですtorch.jit.ScriptModule。通常の PyTorch モデルを TorchScript モデルに変換するには、計算グラフをエクスポートする方法として、トレースとスクリプトの 2 つがあります。torch.onnx.export通常の PyTorch モデル ( ) が渡された場合torch.nn.Module、モデルはデフォルトでトレースメソッドを使用してエクスポートされます。このプロセスを次の図に示します。

最初のチュートリアルの知識を思い出してください。追跡メソッドはモデルを実際に 1 回実行することによってのみモデルの静的グラフをエクスポートできます。つまり、モデル内の制御フロー (ループなど) を識別できません。記録メソッドは正しくエクスポートできます。すべての制御フローを記録します。これら 2 つの変換方法の違いを確認するために、次のコードを例に挙げてみましょう。

import torch 
 
class Model(torch.nn.Module): 
    def __init__(self, n): 
        super().__init__() 
        self.n = n 
        self.conv = torch.nn.Conv2d(3, 3, 3) 
 
    def forward(self, x): 
        for i in range(self.n): 
            x = self.conv(x) 
        return x 
 
 
models = [Model(2), Model(3)] 
model_names = ['model_2', 'model_3'] 
 
for model, model_name in zip(models, model_names): 
    dummy_input = torch.rand(1, 3, 10, 10) 
    dummy_output = model(dummy_input) 
    model_trace = torch.jit.trace(model, dummy_input) 
    model_script = torch.jit.script(model) 
 
    # 跟踪法与直接 torch.onnx.export(model, ...)等价 
    torch.onnx.export(model_trace, dummy_input, f'{model_name}_trace.onnx', example_outputs=dummy_output) 
    # 记录法必须先调用 torch.jit.sciprt 
    torch.onnx.export(model_script, dummy_input, f'{model_name}_script.onnx', example_outputs=dummy_output)

このコードでは、ループを含むモデルを定義し、モデルはn入力テンソルがパラメーターを通じて畳み込まれる回数を制御します。その後、それぞれとn=2の模型を作りましたn=3。これら 2 つのモデルをそれぞれ追跡方法と記録方法でエクスポートします。ここでの 2 つのモデル ( 、 ) は TorchScript モデルであるため、関数はモデルを再度実行する必要がないことに
注意してください。(モデルが追跡メソッドを使用して取得された場合、モデルは実行中に 1 回実行されます。また、記録メソッドがエクスポートに使用された場合、モデルを実際に実行する必要はありません) パラメーター内の and` は、単にデータを取得するためのものです。入力および出力テンソルのタイプと形状。上記のコードを実行すると、Netron で取得した 4 つの onnx ファイルが視覚化されます。model_tracemodel_scriptexporttorch.jit.tracedummy_inputdummy_output

まず、追跡方法によって取得された ONNX モデル構造を見てください。nONNX モデルの構造はモデルごとに異なることがわかります。

この記録方法を使用すると、最終的な ONNX モデルは Loop ノードを使用してループを表します。このように、異なるであっても n、ONNX モデルは同じ構造になります。

この記事で使用されている PyTorch のバージョンは 1.8.2 です。フィードバックによると、PyTorch の他のバージョンでは異なる結果が得られる可能性があります。

推論エンジンは静的グラフをより適切にサポートするため、通常、モデルをデプロイするときに PyTorch モデルを TorchScript モデルに明示的に変換する必要はなく、 Trace を使用して PyTorch モデルを直接エクスポートできます torch.onnx.export 。この部分の知識を理解すると、主に、モデル変換エラーが報告されたときに、PyTorch から TorchScript への段階で問題が発生するかどうかをより正確に特定できます。

パラメータの説明

変換関数の原理を理解した後、関数の主なパラメータの機能を詳しく紹介します。各パラメータの設定方法をすべて列挙するのではなく、アプリケーションの観点から、さまざまなモデル展開シナリオで各パラメータをどのように設定すべきかを主に紹介します。この関数の詳細な API ドキュメントについては、 torch.onnx を参照してください。 PyTorch 1.11.0 ドキュメントは、ファイル内で次のように定義されています
torch.onnx.export 。 torch.onnx.__init__.py

def export(model, args, f, export_params=True, verbose=False, training=TrainingMode.EVAL, 
           input_names=None, output_names=None, aten=False, export_raw_ir=False, 
           operator_export_type=None, opset_version=None, _retain_param_name=True, 
           do_constant_folding=True, example_outputs=None, strip_doc_string=True, 
           dynamic_axes=None, keep_initializers_as_inputs=None, custom_opsets=None, 
           enable_onnx_checker=True, use_external_data_format=False):

最初の 3 つの必須パラメータは、モデル、モデル入力、エクスポートされた onnx ファイルの名前です。これらのパラメータについてはすでによく知っています。後で、一般的に使用されるオプションのパラメーターのいくつかに焦点を当てましょう。

エクスポートパラメータ

モデルの重みをモデルに保存するかどうか。一般に、中間表現にはモデル構造とモデル重みという 2 種類の情報が含まれており、これら 2 種類の情報は同じファイルに保存することも、別のファイルに保存することもできます。ONNX は、同じファイルを使用してレコードモデルの構造と重みを表します。
通常、デプロイ時にはこのパラメータをデフォルトで True に設定します。onnx ファイルをデプロイメントのためではなく、異なるフレームワーク間 (PyTorch から Tensorflow へなど) にモデルを転送するために使用する場合、このパラメーターを False に設定できます。

入力名、出力名

入力テンソルと出力テンソルの名前を設定します。設定されていない場合は、単純な名前 (番号など) が自動的に割り当てられます。
ONNX モデルの入力テンソルと出力テンソルにはそれぞれ名前があります。多くの推論エンジンが ONNX ファイルを実行する場合、「名前とテンソル値」データペアの形式でデータを入力し、出力テンソルの名前に従って出力データを取得する必要があります。テンソル関連の設定 (動的ディメンションの追加など) を行う場合は、テンソルの名前も知っておく必要があります。
実際のデプロイメントパイプラインでは、入力テンソルと出力テンソルの名前を設定し、ONNX と推論エンジンが同じ名前のセットを使用するようにする必要があります。

opset_version

変換時に参照する ONNX オペレータセットのバージョン。デフォルトは 9 です。PyTorchとONNXの演算子の対応については後ほど詳しく紹介します。

動的軸

入力テンソルと出力テンソルのどの次元が動的であるかを指定します。
効率を追求するために、ONNX は、演算に関与するすべてのテンソルが静的である (テンソルの形状が変化しない) ことをデフォルトとしています。しかし、実際のアプリケーションでは、モデルの入力テンソルが動的であること、特に形状制限のない完全畳み込みモデルが望ましいと考えられます。したがって、入力テンソルと出力テンソルのどの次元がサイズ可変であるかを明示的に示す必要があります。
設定例を見てみましょうdynamic_axes。

import torch 
 
class Model(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
        self.conv = torch.nn.Conv2d(3, 3, 3) 
 
    def forward(self, x): 
        x = self.conv(x) 
        return x 
 
 
model = Model() 
dummy_input = torch.rand(1, 3, 10, 10) 
model_names = ['model_static.onnx',  
'model_dynamic_0.onnx',  
'model_dynamic_23.onnx'] 
 
dynamic_axes_0 = { 
    'in' : [0], 
    'out' : [0] 
} 
dynamic_axes_23 = { 
    'in' : [2, 3], 
    'out' : [2, 3] 
} 
 
torch.onnx.export(model, dummy_input, model_names[0],  
input_names=['in'], output_names=['out']) 
torch.onnx.export(model, dummy_input, model_names[1],  
input_names=['in'], output_names=['out'], dynamic_axes=dynamic_axes_0) 
torch.onnx.export(model, dummy_input, model_names[2],  
input_names=['in'], output_names=['out'], dynamic_axes=dynamic_axes_23)

まず、動的次元のないモデル、0 次元のダイナミクス、および 2 次元と 3 次元のダイナミクスを含む 3 つの ONNX モデルをエクスポートします。
このコードでは、次のように動的ディメンションをリストで表します。

dynamic_axes_0 = { 
    'in' : [0], 
    'out' : [0] 
}

ONNX では各動的ディメンションに名前が必要であるため、このように記述すると UserWarning が発生し、リストを通じて動的ディメンションを設定するとシステムが自動的に名前を割り当てることを警告します。動的ディメンション名を明示的に追加する 1 つの方法は次のとおりです。

dynamic_axes_0 = { 
    'in' : {0: 'batch'}, 
    'out' : {0: 'batch'} 
}

このコードでは動的ディメンションに対する操作はこれ以上ないため、リストを使用して動的ディメンションを指定するだけです。
その後、次のコードを使用して、動的ディメンションの役割を見てみましょう。

import onnxruntime 
import numpy as np 
 
origin_tensor = np.random.rand(1, 3, 10, 10).astype(np.float32) 
mult_batch_tensor = np.random.rand(2, 3, 10, 10).astype(np.float32) 
big_tensor = np.random.rand(1, 3, 20, 20).astype(np.float32) 
 
inputs = [origin_tensor, mult_batch_tensor, big_tensor] 
exceptions = dict() 
 
for model_name in model_names: 
    for i, input in enumerate(inputs): 
        try: 
            ort_session = onnxruntime.InferenceSession(model_name) 
            ort_inputs = {'in': input} 
            ort_session.run(['out'], ort_inputs) 
        except Exception as e: 
            exceptions[(i, model_name)] = e 
            print(f'Input[{i}] on model {model_name} error.') 
        else: 
            print(f'Input[{i}] on model {model_name} succeed.')

(1, 3, 10, 10)モデルから計算グラフをエクスポートするときに、形状のテンソルを使用します。ここで、形状を入力として使用し(1, 3, 10, 10), (2, 3, 10, 10), (1, 3, 20, 20)、ONNX ランタイムでこれらのモデルを実行し、どのような状況でエラーが報告されるかを確認し、対応するエラー情報を保存してみましょう。結果の出力は次のようになります。

Input[0] on model model_static.onnx succeed. 
Input[1] on model model_static.onnx error. 
Input[2] on model model_static.onnx error. 
Input[0] on model model_dynamic_0.onnx succeed. 
Input[1] on model model_dynamic_0.onnx succeed. 
Input[2] on model model_dynamic_0.onnx error. 
Input[0] on model model_dynamic_23.onnx succeed. 
Input[1] on model model_dynamic_23.onnx error. 
Input[2] on model model_dynamic_23.onnx succeed.

(1, 3, 10, 10)どの機種でも同じ形状の入力を間違えていないことが分かります。バッチ (0 次元) または長さと幅 (2 次元と 3 次元) が異なる入力の場合、対応する動的次元が設定されるまでエラーは発生しません。どの寸法が間違っているかは、エラーメッセージで確認できます。たとえば、次のコードを使用すると、次のエラーメッセージを表示できinput[1]ますmodel_static.onnx。

print(exceptions[(1, 'model_static.onnx')]) 
 
# output 
# [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Got invalid dimensions for input: in for the following indices index: 0 Got: 2 Expected: 1 Please fix either the inputs or the model.

このエラーメッセージは、in指定された入力の 0 番目の次元が一致しないことを示しています。本来、この次元の長さは 1 である必要がありますが、入力は 2 です。実際のデプロイメントで同様のエラーが発生した場合は、動的ディメンションを設定することで問題を解決できます。

提案を使用する

torch.onnx.exportこれまでの知識を学習することで、関数の部分実現原理とパラメータ設定方法を基本的に習得し、単純なモデルの変換を完了するのに十分です。しかし、実際のアプリケーションでは、この機能を使用すると多くの落とし穴に遭遇することになります。ここでは、モデル展開チームが実際の戦闘で蓄積した経験を共有します。

ONNX で変換したときにモデルの動作が異なるようにする

場合によっては、モデルが ONNX にエクスポートされるときに、モデルにいくつかの異なる動作をさせる必要があります。モデルには、PyTorch で直接推論されるときの 1 つのロジックセットと、エクスポートされた ONNX モデル内の別のロジックセットがあります。たとえば、モデルに後処理ロジックを組み込んで、モデルの実行以外のコードを簡素化できます。torch.onnx.is_in_onnx_export()このタスクを達成するには、関数は torch.onnx.export()実行中にのみ true になります。以下に例を示します。

import torch 
 
class Model(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
        self.conv = torch.nn.Conv2d(3, 3, 3) 
 
    def forward(self, x): 
        x = self.conv(x) 
        if torch.onnx.is_in_onnx_export(): 
            x = torch.clip(x, 0, 1) 
        return x

ここでは、モデルがエクスポートされるときに出力テンソルの値を [0, 1] にのみ制限します。これを使用すると、 is_in_onnx_exportモデルのデプロイメントに関連するロジックをコードに簡単に追加できます。ただし、これらのコードは、モデルのトレーニングのみを気にする開発者やユーザーにとって非常に不親切で、突然のデプロイメントロジックによりコード全体の可読性が低下します。同時に、is_in_onnx_export展開ロジックを追加する必要があるすべての場所に「パッチ」を適用することしかできないため、統合管理を実行することが困難になります。これらの問題を回避するために MMDeploy の書き換えメカニズムを使用する方法を後で紹介します。

割り込みテンソルを使用して追跡された操作

PyTorch から ONNX への追跡エクスポート方法は万能薬ですか? モデル内でいくつかの「すぐに使用できる」操作を実行すると、追跡メソッドは入力に依存するいくつかの中間結果を定数に変換するため、エクスポートされた ONNX モデルは元のモデルとは異なります。この「トレース中断」の原因の例を次に示します。

class Model(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
 
    def forward(self, x): 
        x = x * x[0].item() 
        return x, torch.Tensor([i for i in x]) 
 
model = Model()       
dummy_input = torch.rand(10) 
torch.onnx.export(model, dummy_input, 'a.onnx')

このモデルをエクスポートしようとすると、変換されたモデルが正しくない可能性があることを知らせる警告が大量に表示されます。.item()このモデルでは、トーチのテンソルを使用して通常の Python 変数に変換し、トーチテンソルをトラバースして、リストを含む新しいトーチテンソルを作成しようとしているのは不思議ではありません。テンソルと通常の変数の変換を含むこれらのロジックにより、最終的な ONNX モデルが不正確になります。
一方、この特性を使用して、正確性を保証するという前提の下で、モデルの中間結果が一定になるように命令することもできます。この手法は、モデルを静的にする、つまりモデル内のすべてのテンソル形状が一定になるようにするためによく使用されます。今後のチュートリアルでは、展開例でこれらの「高度な」操作を詳しく説明します。

入力としてテンソルを使用する (PyTorch バージョン < 1.9.0)

最初のチュートリアルで示したように、古い (< 1.9.0) PyTorch は、 torch.onnx.export()Python 値をモデルにフィードするときにエラーをスローします。互換性を確保するために、モデルを変換するときにモデル入力として tensor を使用することをお勧めします。

PyTorch の ONNX オペレーターサポート

呼び出しメソッドが正しいことを確認した後torch.onnx.export()、PyTorch を ONNX に変換するときに最も考えられる問題は、オペレーターに互換性がないことです。ここでは、エラーが発生したときにエラーをより適切に分類できるように、PyTorch オペレーターが ONNX で互換性があるかどうかを判断する方法を紹介します。演算子の具体的な追加方法は、後の記事で紹介します。
通常のモデルを変換する場合torch.nn.Module、PyTorch はトラッキングメソッドを使用して順推論を実行し、検出された演算子を計算グラフに統合します。一方、PyTorch は検出された各演算子を ONNX 演算子に変換します。この翻訳プロセス中に、次の状況が発生する可能性があります。

この演算子は、ONNX 演算子に 1 対 1 で変換できます。
この演算子には ONNX に直接対応する演算子がないため、1 つ以上の ONNX 演算子に変換されます。
オペレーターが ONNX への変換ルールを定義していないため、エラーが報告されます。

では、PyTorch 演算子と ONNX 演算子の対応関係を確認するにはどうすればよいでしょうか? PyTorch 演算子は ONNX に合わせて配置されているため、ここではまず ONNX 演算子の定義を確認し、次に PyTorch によって定義された演算子のマッピング関係を確認します。

ONNX オペレーターのドキュメント

ONNX オペレーターの定義は、オペレーターの公式ドキュメントで参照できます。この文書は非常に重要です。ONNX オペレーターに関連する問題が発生した場合は、この文書を「参照」する必要があります。

この文書の最も重要な部分の冒頭にあるオペレータ変更表。テーブルの最初の列はオペレーター名で、2 番目の列はオペレーターが変更したオペレーターセットのバージョン番号です。これは、前述したtorch.onnx.exportオペレーターopset_versionセットのバージョン番号です。オペレーターの最初の変更のバージョン番号を表示することで、オペレーターがどのバージョンからサポートされているかを知ることができます。オペレーターの最初の変更レコード以下を表示することで、オペレーターセットの現在のバージョンを知ることができますopset_version。の演算子の定義規則。

表内のリンクをクリックすると、入出力パラメータの仕様や演算子の使用例を確認できます。たとえば、上図は ONNX における Relu の定義ルールであり、この定義は、Relu が 1 つの入力と 1 つの入力を持つ必要があり、入力と出力が同じ型 (両方ともテンソル) であることを示しています。

PyTorch から ONNX オペレーターへのマッピング

PyTorch では、次の図に示すように、 ONNX に関連するすべての定義がtorch.onnxディレクトリに配置されます。

このうち、symbolic_opset{n}.py(シンボルテーブルファイル)は、PyTorchがONNXオペレータセットのn番目のバージョンをサポートする際に新たに追加された内容を指します。前に述べたように、バイキュービック補間はバージョン 11 でサポートされています。これを例として、演算子のマッピングを見つける方法を見てみましょう。
まず、検索機能を使用してtorch/onnxフォルダー内で「bicubic」を検索すると、この補間が第 11 バージョンの定義ファイルで見つかります。

その後、コードの呼び出しロジックに従って、一番下の ONNX マッピング関数に段階的にジャンプします。

upsample_bicubic2d = _interpolate("upsample_bicubic2d", 4, "cubic") 
 
-> 
 
def _interpolate(name, dim, interpolate_mode): 
    return sym_help._interpolate_helper(name, dim, interpolate_mode) 
 
-> 
 
def _interpolate_helper(name, dim, interpolate_mode): 
    def symbolic_fn(g, input, output_size, *args): 
        ... 
 
    return symbolic_fn

最後に、ではsymbolic_fn、補間演算子が複数の ONNX 演算子にどのようにマップされるかを確認できます。その中で、それぞれがg.opONNX の定義です。たとえば、 Resize 演算子は次のように記述されます。

return g.op("Resize", 
                input, 
                empty_roi, 
                empty_scales, 
                output_size, 
                coordinate_transformation_mode_s=coordinate_transformation_mode, 
                cubic_coeff_a_f=-0.75,  # only valid when mode="cubic" 
                mode_s=interpolate_mode,  # nearest, linear, or cubic 
                nearest_mode_s="floor")  # only valid when mode="nearest"

前述のONNX オペレーターのドキュメントで Resize オペレーターの定義を調べると、各パラメーターの意味を知ることができます。同様の方法を使用して、他の ONNX オペレーターのパラメーターの意味をクエリし、PyTorch のパラメーターが各 ONNX オペレーターにどのように段階的に渡されるかを知ることができます。
PyTorch と ONNX の間の関係をクエリする方法を習得したら、実際のアプリケーションにバージョン番号を事前設定し torch.onnx.export()、opset_version問題が発生したときに対応する PyTorch シンボルテーブルファイルを確認できます。演算子が存在しない場合、または演算子のマッピング関係が要件を満たしていない場合は、他の演算子でそれをバイパスするか、演算子をカスタマイズする必要がある場合があります。

要約する

このチュートリアルでは、PyTorch を ONNX に変換する原理を体系的に紹介しました。まず、最も頻繁に使用される torch.onnx.export 関数の説明に重点を置き、次に PyTorch の ONNX オペレーターのサポートをクエリするメソッドを示しました。この記事を通じて、新しい演算子を追加する必要がないほとんどの ONNX モデルを正常に変換でき、演算子の問題が発生したときに問題の原因を効果的に特定できることを願っています。具体的には、この記事を読んだ後、次の知識を理解する必要があります。

制御文を含む計算グラフをエクスポートする場合のトレース方法と記録方法の違いは何ですか。
torch.onnx.export()での設定方法です input_names, output_names, dynamic_axes。
torch.onnx.is_in_onnx_export()ONNX に変換したときにモデルの動作を変えるために使用します。
ONNX オペレーターのドキュメント ( https://github.com/onnx/onnx/blob/main/docs/Operators.md )をクエリする方法。
特定の ONNX バージョンの新機能に対する PyTorch のサポートをクエリする方法。
PyTorch が特定の ONNX オペレーターをサポートしているかどうかを判断する方法、およびサポートされている方法は何ですか。

今回紹介した知識は比較的抽象的ですが、ちょっと「水っぽい」と思いませんか？それは問題ではありません。次のチュートリアルでは、PyTorch のオペレーターサポートをコードサンプルの形で ONNX に追加するさまざまな方法を紹介し、PyTorch を ONNX に移行するすべての人にとっての障害をさらに取り除きます。乞うご期待！

練習練習

Asinh オペレーターは、9 番目の ONNX オペレーターセットに表示されます。PyTorch はバージョン 9 のシンボルテーブルファイルでこの演算子をどのようにサポートしますか?
BitShift オペレーターは、11 番目の ONNX オペレーターセットに登場しました。PyTorch はバージョン 11 のシンボルテーブルファイルでこの演算子をどのようにサポートしますか?
最初のチュートリアルで、PyTorch (オペレーターセット #11 以降) は補間における動的なスケーリング係数の設定をサポートしていないと述べました。この係数は、 Resize オペレーターのマッピング関係のどのパラメーターに対応しますかtorch.onnx.symbolic_helper._interpolate_helper? symbolic_fnこのパラメータをどのように変更したのでしょうか?