Udostępnij przez


Konwertowanie modelu PyTorch na format 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 DataClassifier.py w programie Visual Studio, przed funkcją główną.
#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') 

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.

  1. 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__": 
    num_epochs = 10 
    train(num_epochs) 
    print('Finished Training\n') 
    test() 
    test_species() 
    convert() 
  1. 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.

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

Uwaga / Notatka

Chcesz dowiedzieć się więcej? Zapoznaj się z samouczkiem PyTorch dotyczącym eksportowania mdoelu.

Eksploruj swój model.

  1. Otwórz plik modelu Network.onnx w programie Neutron.

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

Właściwości modelu ONNX

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. Model został zbudowany w taki sposób, że etykiety są przedstawione przez 3 liczby, z których każda jest przypisana do określonego typu kwiatu irysa.

Etykieta 1 etykieta 2 Etykieta 3
0 1 2
Irys-setosa Iris-versicolor (kosaciec różnobarwny) Irys-wirginijka

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.