PyTorch トレーニング モデルを ONNX に変換する

このチュートリアルの前の段階では、PyTorch を使用して機械学習モデルを作成しました。 ただし、そのモデルは、.pth ファイルです。 Windows ML アプリと統合するには、モデルを ONNX 形式に変換する必要があります。

モデルのエクスポート

モデルをエクスポートするには、torch.onnx.export() 関数を使用します。 この関数によってモデルを実行し、出力を計算するために使用される演算子のトレースを記録します。

1. Visual Studio で、次のコードを PyTorchTraining.py ファイルの main 関数の上にコピーします。

import torch.onnx 

#Function to Convert to ONNX 
def Convert_ONNX(): 

    # set the model to inference mode 
    model.eval() 

    # Let's create a dummy input tensor  
    dummy_input = torch.randn(1, input_size, requires_grad=True)  

    # Export the model   
    torch.onnx.export(model,         # model being run 
         dummy_input,       # model input (or a tuple for multiple inputs) 
         "ImageClassifier.onnx",       # where to save the model  
         export_params=True,  # store the trained parameter weights inside the model file 
         opset_version=10,    # the ONNX version to export the model to 
         do_constant_folding=True,  # whether to execute constant folding for optimization 
         input_names = ['modelInput'],   # the model's input names 
         output_names = ['modelOutput'], # the model's output names 
         dynamic_axes={'modelInput' : {0 : 'batch_size'},    # variable length axes 
                                'modelOutput' : {0 : 'batch_size'}}) 
    print(" ") 
    print('Model has been converted to ONNX') 

モデルをエクスポートする前に model.eval() または model.train(False) を呼び出すことが重要です。これにより、モデルが推論モードに設定されます。 dropout や batchnorm などの演算子は、推論モードとトレーニング モードでは動作が異なります。

2. ONNX への変換を実行するには、変換関数の呼び出しを main 関数に追加します。 モデルを再度トレーニングする必要はないので、実行する必要がなくなったいくつかの関数をコメント アウトします。 main 関数は次のようになります。

if __name__ == "__main__": 

    # Let's build our model 
    #train(5) 
    #print('Finished Training') 

    # Test which classes performed well 
    #testAccuracy() 

    # Let's load the model we just created and test the accuracy per label 
    model = Network() 
    path = "myFirstModel.pth" 
    model.load_state_dict(torch.load(path)) 

    # Test with batch of images 
    #testBatch() 
    # Test how the classes performed 
    #testClassess() 
 
    # Conversion to ONNX 
    Convert_ONNX() 

3. ツール バーの [Start Debugging] ボタンを選択するか F5 キーを押して、プロジェクトを再度実行します。 モデルを再度トレーニングする必要はありません。プロジェクト フォルダーから既存のモデルを読み込むだけです。

出力は次のようになります。

プロジェクトの場所に移動し、.pth モデルの横にある ONNX モデルを探します。

Note

もっと詳しく知りたいですか? モデルのエクスポートに関する PyTorch チュートリアルを参照してください。

モデルを探索する

1. Netron を使用して ImageClassifier.onnx モデル ファイルを開きます。

2. "データ" ノードを選択して、モデルのプロパティを開きます。

ご覧のように、このモデルでは、入力として 32 ビットのテンソル (多次元配列) 浮動小数点数オブジェクトが必要であり、出力として Tensor float が返されます。 出力配列には、すべてのラベルの確率が含まれます。 このモデルの構築方法では、ラベルは 10 個の数値で表され、すべての数値がオブジェクトの 10 個のクラスを表します。

ラベル 0	ラベル 1	ラベル 2	ラベル 3	ラベル 4	ラベル 5	ラベル 6	ラベル 7	ラベル 8	ラベル 9
0	1	2	3	4	5	6	7	8	9
航空機	car	bird	cat	deer	犬	frog	horse	ship	truck

Windows ML アプリで正しい予測を表示するには、これらの値を抽出する必要があります。

次のステップ

モデルをデプロイする準備ができました。 次に、メイン イベントである Windows アプリケーションのビルドと Windows デバイス上でのローカルの実行を行いましょう。