Share via


Converter o modelo PyTorch no formato 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 DataClassifier.py no Visual Studio, acima da função principal.
#Function to Convert to ONNX 
def convert(): 

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

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

    # Export the model   
    torch.onnx.export(model,         # model being run 
         dummy_input,       # model input (or a tuple for multiple inputs) 
         "Network.onnx",       # where to save the model  
         export_params=True,  # store the trained parameter weights inside the model file 
         opset_version=11,    # the ONNX version to export the model to 
         do_constant_folding=True,  # whether to execute constant folding for optimization 
         input_names = ['input'],   # the model's input names 
         output_names = ['output'], # the model's output names 
         dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes 
                                'output' : {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__": 
    num_epochs = 10 
    train(num_epochs) 
    print('Finished Training\n') 
    test() 
    test_species() 
    convert() 
  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.

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

Observação

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

Explorar o modelo.

  1. Abra o arquivo de modelo Network.onnx com o Neutron.

  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. O modelo foi criado de modo que os rótulos fossem representados por três números, cada um associado a uma espécie específica de flor de íris.

Rótulo 1 rótulo 2 Rótulo 3
0 1 2
{i>Iris setosa {i>Iris versicolor {i>Iris virginica

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.