Sdílet prostřednictvím


Převod modelu PyTorch na formát ONNX

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 DataClassifier.py souboru ve Visual Studio, před hlavní funkcí.
#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') 

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í.

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

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 mdoelu.

Prozkoumejte svůj model.

  1. Otevřete soubor modelu pomocí Neutronu.

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

Vlastnosti modelu ONNX

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, jakým jste model vytvořili, jsou popisky reprezentovány 3 čísly, z nichž každá je přidružená ke konkrétnímu typu duhovky.

Popisek 1 popisek 2 Popisek 3
0 1 2
Kosatec setosa Iris-versicolor Iris virginica

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.