Converta o seu modelo PyTorch para o formato 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 DataClassifier.py arquivo no Visual Studio, acima de sua 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 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__": 
    num_epochs = 10 
    train(num_epochs) 
    print('Finished Training\n') 
    test() 
    test_species() 
    convert() 

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.

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 o Network.onnx Neutron.

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 forma como você construiu o modelo, os rótulos são representados por 3 números, cada um associado a um tipo específico de flor de íris.

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.