Conversión del modelo de PyTorch al formato ONNX

En la fase anterior de este tutorial usó PyTorch para crear el modelo de aprendizaje automático. Sin embargo, ese modelo es un archivo .pth. Para poder integrarlo con la aplicación Windows ML, deberá convertir el modelo al formato ONNX.

Exportación del modelo

Para exportar un modelo, debe usar la función torch.onnx.export(). Esta función ejecuta el modelo y registra un seguimiento de los operadores que se usan para calcular las salidas.

1. Copie el código siguiente en el archivo DataClassifier.py de Visual Studio, encima de la función principal.

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

Es importante llamar a model.eval() o a model.train(False) antes de exportar el modelo, ya que esto establece el modelo en el modo de inferencia. Esto es necesario porque los operadores como dropout o batchnorm se comportan de forma diferente en el modo de inferencia y entrenamiento.

2. Para ejecutar la conversión a ONNX, agregue una llamada a la función de conversión a la función principal. No es necesario volver a entrenar el modelo, por lo que comentaremos algunas funciones que ya no necesitamos ejecutar. La función principal será la siguiente.

if __name__ == "__main__": 
    num_epochs = 10 
    train(num_epochs) 
    print('Finished Training\n') 
    test() 
    test_species() 
    convert() 

3. Vuelva a ejecutar el proyecto seleccionando el botón Start Debugging de la barra de herramientas o presionando F5. No es necesario volver a entrenar el modelo; simplemente, cargue el modelo existente desde la carpeta del proyecto.

Vaya a la ubicación del proyecto y busque el modelo ONNX junto al modelo .pth.

Nota:

¿Quiere saber más sobre el tema? Revise el tutorial de PyTorch sobre la exportación de un modelo.

Explore el modelo.

1. Abra el archivo de modelo Network.onnx con Neutron.

2. Seleccione el nodo de datos para abrir las propiedades del modelo.

Como puede ver, el modelo requiere un objeto de tensor flotante de 32 bits (esto es, una matriz multidimensional) como entrada y devuelve un valor "float" de Tensor como salida. La matriz de salida incluirá la probabilidad de cada etiqueta. Por la forma en que creó el modelo, las etiquetas se representan mediante tres números, cada uno asociado a un tipo específico de flor Iris.

Etiqueta 1	label 2	Etiqueta 3
0	1	2
Iris setosa	Iris versicolor	Iris virginica

Deberá extraer estos valores para mostrar la predicción correcta con aplicación la Windows ML.

Pasos siguientes

El modelo está listo para implementarse. A continuación, para el evento principal, tendrá que compilar una aplicación de Windows y ejecutarla localmente en el dispositivo Windows.