Konwertowanie modelu trenowania PyTorch na ONNX

W poprzednim etapie tego samouczka użyliśmy usługi PyTorch do utworzenia naszego modelu uczenia maszynowego. Jednak ten model jest plikiem .pth . Aby móc ją zintegrować z aplikacją uczenia maszynowego systemu Windows, należy przekonwertować model na format ONNX.

Eksportowanie modelu

Aby wyeksportować model, użyjesz torch.onnx.export() funkcji . Ta funkcja wykonuje model i rejestruje ślad operatorów używanych do obliczania danych wyjściowych.

1. Skopiuj następujący kod do pliku PyTorchTraining.py w programie Visual Studio, przed funkcją główną.

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

Ważne jest wywołanie model.eval() lub model.train(False) przed wyeksportowaniem modelu, ponieważ powoduje to ustawienie modelu na tryb wnioskowania. Jest to potrzebne, ponieważ operatory takie jak dropout lub batchnorm zachowują się inaczej w trybie wnioskowania i treningu.

2. Aby uruchomić konwersję na ONNX, w głównej funkcji dodaj wywołanie funkcji konwersji. Nie musisz ponownie trenować modelu, dlatego oznaczymy niektóre funkcje, których nie musimy już uruchamiać. Główna funkcja będzie następująca.

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. Ponownie uruchom projekt, wybierając Start Debugging przycisk na pasku narzędzi lub naciskając F5. Nie ma potrzeby ponownego trenowania modelu— wystarczy załadować istniejący model z folderu projektu.

Dane wyjściowe będą wyglądać następująco.

Przejdź do lokalizacji projektu i znajdź model ONNX obok .pth modelu.

Eksploruj swój model.

1. Otwórz plik modelu w Netron.

2. Wybierz węzeł danych , aby otworzyć właściwości modelu.

Jak widać, model wymaga obiektu zmiennoprzecinkowego 32-bitowego (tablicy wielowymiarowej) jako danych wejściowych i zwraca zmiennoprzecinkowy Tensor jako dane wyjściowe. Tablica wyjściowa będzie zawierać prawdopodobieństwo dla każdej etykiety. Sposób skompilowania modelu, etykiety są reprezentowane przez 10 liczb, a każda liczba reprezentuje dziesięć klas obiektów.

Należy wyodrębnić te wartości, aby pokazać poprawne przewidywanie za pomocą aplikacji Windows ML.

Dalsze kroki

Nasz model jest gotowy do wdrożenia. Następnie dla zdarzenia głównego — skompilujmy aplikację systemu Windows i uruchomimy ją lokalnie na urządzeniu z systemem Windows.