Converter o modelo de treinamento PyTorch em ONNX

Na etapa anterior deste tutorial, usamos PyTorch para criar o modelo de machine learning. No entanto, esse modelo é um arquivo .pth. Para integrá-lo com o aplicativo do Windows ML, você vai precisar converter o modelo no formato ONNX.

Exportar o modelo

Para exportar um modelo, você vai usar a função torch.onnx.export(). Essa função executa o modelo e registra o rastreamento de quais operadores são usados para computar as saídas.

  1. Copie o código a seguir no arquivo PyTorchTraining.py no Visual Studio, acima da função principal.
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') 

É importante chamar model.eval() ou model.train(False) antes de exportar o modelo, pois isso o define para o modo de inferência. Isso é necessário, pois operadores como dropout ou batchnorm se comportam de maneiras diferentes no modo de inferência e de treinamento.

  1. Para executar a conversão em ONNX, adicione uma chamada à função de conversão para a função principal. Você não precisa treinar o modelo novamente, então vamos transformar em comentários algumas funções que não precisam mais ser executadas. A sua função principal será conforme demonstrado a seguir.
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() 
  1. Execute o projeto novamente selecionando o botão Start Debugging na barra de ferramentas ou pressionando F5. Não há necessidade de treinar o modelo de novo, basta carregar o modelo existente da pasta do projeto.

A saída será conforme demonstrado a seguir.

ONNX conversion process

Navegue até o local do projeto e localize o modelo ONNX ao lado do modelo .pth.

Nota

Quer aprender mais? Examine o Tutorial do PyTorch sobre como exportar um modelo.

Explorar o modelo.

  1. Abra o arquivo de modelo ImageClassifier.onnx com o Netron.

  2. Selecione o nó de dados para abrir as propriedades do modelo.

ONNX model properties

Como você pode ver, o modelo requer um objeto float de tensor de 32 bits (matriz multidimensional) como uma entrada e retorna um tensor float como uma saída. A matriz de saída inclui a probabilidade de cada rótulo. Da maneira como o modelo foi criado, os rótulos são representados por 10 números e cada número representa as dez classes de objetos.

Rótulo 0 Rótulo 1 Rótulo 2 Rótulo 3 Rótulo 4 Rótulo 5 Rótulo 6 Rótulo 7 Rótulo 8 Rótulo 9
0 1 2 3 4 5 6 7 8 9
avião car bird cat veado cachorro sapo cavalo navio caminhão

Você vai precisar extrair esses valores para mostrar a previsão correta com aplicativo do Windows ML.

Próximas etapas

O seu modelo está pronto para ser implantado. Em seguida, para o evento principal, vamos criar um aplicativo Windows e executá-lo localmente no seu dispositivo Windows.