Dela via


Konvertera PyTorch-modellen till ONNX-format

I föregående steg i den här självstudien använde vi PyTorch för att skapa vår maskininlärningsmodell. Den modellen är dock en .pth fil. För att kunna integrera den med Windows ML-appen måste du konvertera modellen till ONNX-format.

Exportera modellen

Om du vill exportera en modell använder torch.onnx.export() du funktionen. Den här funktionen kör modellen och registrerar en spårning av vilka operatorer som används för att beräkna utdata.

  1. Kopiera följande kod till DataClassifier.py filen i Visual Studio ovanför huvudfunktionen.
#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') 

Det är viktigt att anropa model.eval() eller model.train(False) innan du exporterar modellen, eftersom detta ställer in modellen till slutsatsdragningsläge. Detta behövs eftersom operatorer som dropout eller batchnorm beter sig annorlunda i inferens- och träningsläge.

  1. Om du vill köra konverteringen till ONNX lägger du till ett anrop till konverteringsfunktionen i huvudfunktionen. Du behöver inte träna modellen igen, så vi kommenterar ut några funktioner som vi inte längre behöver köra. Huvudfunktionen är följande.
if __name__ == "__main__": 
    num_epochs = 10 
    train(num_epochs) 
    print('Finished Training\n') 
    test() 
    test_species() 
    convert() 
  1. Kör projektet igen genom att välja knappen Start Debugging i verktygsfältet eller trycka på F5. Du behöver inte träna modellen igen. Läs bara in den befintliga modellen från projektmappen.

Gå till projektplatsen och leta reda på ONNX-modellen bredvid .pth modellen.

Anmärkning

Vill du veta mer? Granska PyTorch-självstudien om hur du exporterar en mdoel.

Utforska din modell.

  1. Network.onnx Öppna modellfilen med Neutron.

  2. Välj datanoden för att öppna modellegenskaperna.

Egenskaper för ONNX-modell

Som du ser kräver modellen ett 32-bitars tensor-objekt (flerdimensionell matris) som indata och returnerar en Tensor-flyttal som utdata. Utdatamatrisen inkluderar sannolikheten för varje etikett. På det sätt som du skapade modellen representeras etiketterna av 3 tal, var och en associerad med en viss typ av irisblomma.

Etikett 1 etikett 2 Etikett 3
0 1 2
Iris-setosa Iris-versicolor Iris-virginica

Du måste extrahera dessa värden för att visa rätt förutsägelse med Windows ML-appen.

Nästa steg

Vår modell är redo att distribueras. För huvudhändelsen ska vi sedan skapa ett Windows-program och köra det lokalt på din Windows-enhet.