Converta seu modelo de treinamento PyTorch para ONNX

Na etapa anterior deste tutorial, usamos o PyTorch para criar nosso modelo de aprendizado de máquina. No entanto, esse modelo é um .pth arquivo. Para poder integrá-lo ao aplicativo Windows ML, você precisará converter o modelo para o formato ONNX.

Exportar o modelo

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

1. Copie o código a seguir para o PyTorchTraining.py arquivo no Visual Studio, acima de sua 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 define o modelo para o modo de inferência. Isso é necessário, uma vez que os operadores como dropout ou batchnorm se comportam de forma diferente no modo de inferência e modo de treino.

2. Para executar a conversão para 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 comentar algumas funções que não precisamos mais executar. A sua principal função será a seguinte.

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. Execute o projeto novamente selecionando o Start Debugging botão na barra de ferramentas ou pressionando F5. Não há necessidade de treinar o modelo novamente, basta carregar o modelo existente da pasta do projeto.

A saída será a seguinte.

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

Explore o seu modelo.

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

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

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

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

Próximas Etapas

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