Převod trénovacího modelu PyTorch na ONNX

Poznámka:

Pro větší funkčnost lze PyTorch použít také s DirectML ve Windows.

V předchozí fázi tohoto kurzu jsme použili PyTorch k vytvoření našeho modelu strojového učení. Tento model je ale .pth soubor. Abyste ho mohli integrovat s aplikací Windows ML, budete muset model převést na formát ONNX.

Export modelu

K exportu modelu použijete torch.onnx.export() funkci. Tato funkce spustí model a zaznamená sled operátorů, které jsou použity k výpočtu výstupů.

1. Zkopírujte následující kód do PyTorchTraining.py souboru ve Visual Studio, před hlavní funkcí.

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') 

Je důležité volat model.eval() nebo model.train(False) před exportem modelu, protože tím model nastavíte na režim odvozování. To je potřeba, protože operátory jako dropout nebo batchnorm se chovají jinak v režimu odvozování a v režimu trénování.

2. Chcete-li spustit převod na ONNX, přidejte volání funkce převodu do hlavní funkce. Model už nemusíte trénovat, takže některé funkce, které už nemusíme spouštět, zakomentujeme. Hlavní funkce bude následující.

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. Spusťte projekt znovu výběrem tlačítka na panelu nástrojů Start Debugging nebo stisknutím F5. Model už nemusíte trénovat znovu, stačí načíst stávající model ze složky projektu.

Výstup bude následující.

Přejděte do umístění projektu a vyhledejte model ONNX vedle .pth modelu.

Poznámka:

Zajímá vás další informace? Projděte si kurz PyTorch o exportu modelu.

Prozkoumejte svůj model.

1. ImageClassifier.onnx Otevřete soubor modelu pomocí Netronu.

2. Výběrem datového uzlu otevřete vlastnosti modelu.

Jak můžete vidět, model vyžaduje 32bitový objekt tensoru (multidimenzionální pole) plovoucí jako vstup a vrátí tensor float jako výstup. Výstupní pole bude obsahovat pravděpodobnosti každého štítku. Způsob vytvoření modelu je takový, že popisky jsou reprezentovány 10 čísly a každé číslo představuje jednu z deseti tříd objektů.

Popisek 0	Popisek 1	Popisek 2	Popisek 3	Popisek 4	Popisek 5	Popisek 6	Popisek 7	Popisek 8	Popisek 9
0	1	2	3	4	5	6	7	8	9
letadlo	auto	pták	kočka	jelen	pes	žába	kůň	loď	nákladní auto

Abyste mohli zobrazit správnou předpověď v aplikaci Windows ML, budete muset tyto hodnoty extrahovat.

Další kroky

Náš model je připravený k nasazení. V dalším kroku pro hlavní událost – pojďme sestavit aplikaci pro Windows a spustit ji místně na zařízení s Windows.