Obtención de predicciones con ONNX en modelos de Computer Vision desde AutoML

SE APLICA A: SDK de Python azure-ai-ml v2 (actual)

En este artículo, aprenderá a usar Open Neural Network Exchange (ONNX) para hacer predicciones en modelos de Computer Vision generados con el aprendizaje automático automatizado (AutoML) en Azure Machine Learning.

Para usar ONNX para las predicciones, debe hacer lo siguiente:

  1. Descargue los archivos del modelo de ONNX de una ejecución de entrenamiento de AutoML.
  2. Comprenda las entradas y las salidas de un modelo de ONNX.
  3. Preprocese los datos para que estén en el formato necesario para las imágenes de entrada.
  4. Haga la inferencia con ONNX Runtime for Python.
  5. Visualice las predicciones para las tareas de detección de objeto visual y segmentación de instancias.

ONNX es un estándar abierto para los modelos de aprendizaje automático y aprendizaje profundo. Permite la importación y exportación los de modelos (interoperabilidad) en los marcos de IA populares. Para obtener más información, explore el proyecto GitHub ONNX.

ONNX Runtime es un proyecto de código abierto que admite la inferencia multiplataforma. ONNX Runtime proporciona las API en lenguajes de programación (incluidos Python, C++, C#, C, Java y JavaScript). Puede usar estas API para hacer inferencias en imágenes de entrada. Una vez que tenga el modelo que se exportó al formato ONNX, puede usar estas API en cualquier lenguaje de programación que necesite el proyecto.

En esta guía, aprenderá a usar las API de Python para ONNX Runtime para hacer las predicciones en las imágenes para las tareas de visión populares. Puede usar estos modelos exportados de ONNX entre lenguajes.

Prerrequisitos

Descarga de archivos del modelo de ONNX

Puede descargar archivos del modelo de ONNX de las ejecuciones de AutoML mediante la UI de Estudio de Azure Machine Learning o el SDK de Python de Azure Machine Learning. Se recomienda descargar mediante el SDK con el nombre del experimento y el identificador de ejecución primario.

Azure Machine Learning Studio

En Estudio de Azure Machine Learning, vaya al experimento mediante el hipervínculo al experimento generado en el cuaderno de entrenamiento o mediante la selección del nombre del experimento en la pestaña Experiments (Experimentos) en Assets (Recursos). A continuación, seleccione la mejor ejecución secundaria.

En la mejor ejecución secundaria, vaya a Outputs+logs>train_artifacts. Use el botón Download para descargar manualmente los siguientes archivos:

  • labels.json: archivo que contiene todas las clases o etiquetas del conjunto de datos de entrenamiento.
  • model.onnx: modelo en formato ONNX.

Screenshot that shows selections for downloading O N N X model files.

Guarde los archivos descargados del modelo en un directorio. En el ejemplo de este artículo se usa el directorio ./automl_models.

SDK de Python de Azure Machine Learning

Con el SDK, puede seleccionar la mejor ejecución secundaria (por métrica principal) con el nombre del experimento y el id. de ejecución primario. A continuación, puede descargar los archivos labels.json y model.onnx.

El código siguiente devuelve la mejor ejecución secundaria en función de la métrica principal pertinente.

from azure.identity import DefaultAzureCredential
from azure.ai.ml import MLClient
mlflow_client = MlflowClient()

credential = DefaultAzureCredential()
ml_client = None
try:
    ml_client = MLClient.from_config(credential)
except Exception as ex:
    print(ex)
    # Enter details of your Azure Machine Learning workspace
    subscription_id = ''   
    resource_group = ''  
    workspace_name = ''
    ml_client = MLClient(credential, subscription_id, resource_group, workspace_name)
import mlflow
from mlflow.tracking.client import MlflowClient

# Obtain the tracking URL from MLClient
MLFLOW_TRACKING_URI = ml_client.workspaces.get(
    name=ml_client.workspace_name
).mlflow_tracking_uri

mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)

# Specify the job name
job_name = ''

# Get the parent run
mlflow_parent_run = mlflow_client.get_run(job_name)
best_child_run_id = mlflow_parent_run.data.tags['automl_best_child_run_id']
# get the best child run
best_run = mlflow_client.get_run(best_child_run_id)

Descargue el archivo labels.json que contiene todas las clases y etiquetas del conjunto de datos de entrenamiento.

local_dir = './automl_models'
if not os.path.exists(local_dir):
    os.mkdir(local_dir)

labels_file = mlflow_client.download_artifacts(
    best_run.info.run_id, 'train_artifacts/labels.json', local_dir
)

Descargue el archivo model.onnx.

onnx_model_path = mlflow_client.download_artifacts(
    best_run.info.run_id, 'train_artifacts/model.onnx', local_dir
)

En el caso de la inferencia por lotes para la detección de objetos y la segmentación de instancias mediante modelos ONNX, consulte la sección sobre la generación de modelos para la puntuación por lotes.

Generación de modelos para la puntuación por lotes

De manera predeterminada, AutoML para imágenes admite la puntuación por lotes para la clasificación. Pero los modelos ONNX de detección de objetos y segmentación de instancias no admiten la inferencia por lotes. En el caso de la inferencia por lotes para la detección de objetos y la segmentación de instancias, use el siguiente procedimiento para generar un modelo ONNX para el tamaño de lote necesario. Los modelos generados para un tamaño de lote específico no funcionan para otros tamaños de lote.

Descargue el archivo de entorno de Conda y cree un objeto de entorno que se usará con el trabajo de comando.

#  Download conda file and define the environment

conda_file = mlflow_client.download_artifacts(
    best_run.info.run_id, "outputs/conda_env_v_1_0_0.yml", local_dir
)
from azure.ai.ml.entities import Environment
env = Environment(
    name="automl-images-env-onnx",
    description="environment for automl images ONNX batch model generation",
    image="mcr.microsoft.com/azureml/openmpi4.1.0-cuda11.1-cudnn8-ubuntu18.04",
    conda_file=conda_file,
)

Use los siguientes argumentos específicos del modelo para enviar el script. Para más información sobre los argumentos, consulte los hiperparámetros específicos del modelo y para conocer los nombres de modelos de detección de objetos admitidos, consulte la sección de arquitectura de modelos admitidos.

Para obtener los valores de argumento necesarios para crear el modelo de puntuación por lotes, consulte los scripts de puntuación generados en la carpeta outputs de las ejecuciones de entrenamiento de AutoML. Use los valores de hiperparámetro disponibles en la variable de configuración del modelo dentro del archivo de puntuación para la mejor ejecución secundaria.

Para la clasificación de imágenes de varias clases, el modelo ONNX generado para la mejor ejecución secundaria admite la puntuación por lotes de manera predeterminada. Por lo tanto, no se necesitan argumentos específicos del modelo para este tipo de tarea y puede ir directamente a la sección Carga de las etiquetas y los archivos del modelo ONNX.

Descargue y mantenga el archivo ONNX_batch_model_generator_automl_for_images.py en el directorio actual para enviar el script. Use el siguiente comando de trabajo para enviar el script ONNX_batch_model_generator_automl_for_images.py disponible en el repositorio de GitHub azureml-examples para generar un modelo ONNX de un tamaño de lote específico. En el código siguiente, el entorno del modelo entrenado se usa para enviar este script para generar y guardar el modelo ONNX en el directorio outputs.

Para la clasificación de imágenes de varias clases, el modelo ONNX generado para la mejor ejecución secundaria admite la puntuación por lotes de manera predeterminada. Por lo tanto, no se necesitan argumentos específicos del modelo para este tipo de tarea y puede ir directamente a la sección Carga de las etiquetas y los archivos del modelo ONNX.

Una vez generado el modelo por lotes, descárguelo de Salidas y registros>Salidas manualmente mediante la UI o use el método siguiente:

batch_size = 8  # use the batch size used to generate the model
returned_job_run = mlflow_client.get_run(returned_job.name)

# Download run's artifacts/outputs
onnx_model_path = mlflow_client.download_artifacts(
    returned_job_run.info.run_id, 'outputs/model_'+str(batch_size)+'.onnx', local_dir
)

Después del paso de descarga del modelo, use el paquete de Python de ONNX Runtime para hacer la inferencia mediante el archivo model.onnx. Con fines de demostración, en este artículo se usan los conjuntos de datos de Preparación de conjuntos de datos de imagen para cada tarea de visión.

Hemos entrenado los modelos para todas las tareas de visión con sus respectivos conjuntos de datos para demostrar la inferencia de los modelos de ONNX.

Carga de las etiquetas y los archivos del modelo de ONNX

El siguiente fragmento de código carga labels.json, en el que se ordenan los nombres de clase. Es decir, si el modelo de ONNX predice un id. de etiqueta como 2, corresponde al nombre de etiqueta especificado en el tercer índice del archivo labels.json.

import json
import onnxruntime

labels_file = "automl_models/labels.json"
with open(labels_file) as f:
    classes = json.load(f)
print(classes)
try:
    session = onnxruntime.InferenceSession(onnx_model_path)
    print("ONNX model loaded...")
except Exception as e: 
    print("Error loading ONNX file: ", str(e))

Obtención de los detalles de entrada y salida esperados para un modelo de ONNX

Cuando tenga el modelo, es importante conocer algunos detalles específicos del modelo y específicos de la tarea. Estos detalles incluyen el número de entradas y el número de salidas, la forma o el formato de entrada esperados para preprocesar la imagen y la forma de salida para que conozca las salidas específicas del modelo o específicas de la tarea.

sess_input = session.get_inputs()
sess_output = session.get_outputs()
print(f"No. of inputs : {len(sess_input)}, No. of outputs : {len(sess_output)}")

for idx, input_ in enumerate(range(len(sess_input))):
    input_name = sess_input[input_].name
    input_shape = sess_input[input_].shape
    input_type = sess_input[input_].type
    print(f"{idx} Input name : { input_name }, Input shape : {input_shape}, \
    Input type  : {input_type}")  

for idx, output in enumerate(range(len(sess_output))):
    output_name = sess_output[output].name
    output_shape = sess_output[output].shape
    output_type = sess_output[output].type
    print(f" {idx} Output name : {output_name}, Output shape : {output_shape}, \
    Output type  : {output_type}") 

Formatos de entrada y salida esperados para el modelo de ONNX

Cada modelo de ONNX tiene un conjunto predefinido de formatos de entrada y salida.

En este ejemplo se aplica el modelo entrenado en el conjunto de datos fridgeObjects con 134 imágenes y 4 clases o etiquetas para explicar la inferencia del modelo de ONNX. Para obtener más información sobre cómo entrenar una tarea de clasificación de imágenes, consulte el cuaderno de clasificación de imágenes de varias clases.

Formato de entrada

La entrada es una imagen preprocesada.

Nombre de entrada Forma de entrada Tipo de entrada Descripción
input1 (batch_size, num_channels, height, width) ndarray(float) La entrada es una imagen preprocesada, con la forma (1, 3, 224, 224) para un tamaño de lote de 1 y un alto y ancho de 224. Estos números corresponden a los valores usados para crop_size del ejemplo de entrenamiento.

Formato de salida

La salida es una matriz de logits para todas las clases o etiquetas.

Nombre de salida Forma de salida Tipo de salida Descripción
output1 (batch_size, num_classes) ndarray(float) El modelo devuelve logits (sin softmax). Por ejemplo, para las clases de tamaño de lote 1 y 4, devuelve (1, 4).

Preprocessing (Preprocesamiento)

Siga los siguientes pasos de preprocesamiento para la inferencia del modelo de ONNX:

  1. Convierta la imagen a RGB.
  2. Cambie el tamaño de la imagen a los valores valid_resize_size y valid_resize_size correspondientes a los valores usados en la transformación del conjunto de datos de validación durante el entrenamiento. El valor predeterminado para valid_resize_size es 256.
  3. Recorte en el centro la imagen a height_onnx_crop_size y width_onnx_crop_size. Esto corresponde a valid_crop_size con el valor predeterminado de 224.
  4. Cambio de HxWxC a CxHxW.
  5. Convierta al tipo float.
  6. Normalice con mean = [0.485, 0.456, 0.406] y std = [0.229, 0.224, 0.225] de ImageNet.

Si elige valores distintos para los hiperparámetrosvalid_resize_size y valid_crop_size durante el entrenamiento, se deben usar esos valores.

Obtenga la forma de entrada necesaria para el modelo de ONNX.

batch, channel, height_onnx_crop_size, width_onnx_crop_size = session.get_inputs()[0].shape
batch, channel, height_onnx_crop_size, width_onnx_crop_size

Sin PyTorch

import glob
import numpy as np
from PIL import Image

def preprocess(image, resize_size, crop_size_onnx):
    """Perform pre-processing on raw input image
    
    :param image: raw input image
    :type image: PIL image
    :param resize_size: value to resize the image
    :type image: Int
    :param crop_size_onnx: expected height of an input image in onnx model
    :type crop_size_onnx: Int
    :return: pre-processed image in numpy format
    :rtype: ndarray 1xCxHxW
    """

    image = image.convert('RGB')
    # resize
    image = image.resize((resize_size, resize_size))
    #  center  crop
    left = (resize_size - crop_size_onnx)/2
    top = (resize_size - crop_size_onnx)/2
    right = (resize_size + crop_size_onnx)/2
    bottom = (resize_size + crop_size_onnx)/2
    image = image.crop((left, top, right, bottom))

    np_image = np.array(image)
    # HWC -> CHW
    np_image = np_image.transpose(2, 0, 1) # CxHxW
    # normalize the image
    mean_vec = np.array([0.485, 0.456, 0.406])
    std_vec = np.array([0.229, 0.224, 0.225])
    norm_img_data = np.zeros(np_image.shape).astype('float32')
    for i in range(np_image.shape[0]):
        norm_img_data[i,:,:] = (np_image[i,:,:]/255 - mean_vec[i])/std_vec[i]
             
    np_image = np.expand_dims(norm_img_data, axis=0) # 1xCxHxW
    return np_image

# following code loads only batch_size number of images for demonstrating ONNX inference
# make sure that the data directory has at least batch_size number of images

test_images_path = "automl_models_multi_cls/test_images_dir/*" # replace with path to images
# Select batch size needed
batch_size = 8
# you can modify resize_size based on your trained model
resize_size = 256
# height and width will be the same for classification
crop_size_onnx = height_onnx_crop_size 

image_files = glob.glob(test_images_path)
img_processed_list = []
for i in range(batch_size):
    img = Image.open(image_files[i])
    img_processed_list.append(preprocess(img, resize_size, crop_size_onnx))
    
if len(img_processed_list) > 1:
    img_data = np.concatenate(img_processed_list)
elif len(img_processed_list) == 1:
    img_data = img_processed_list[0]
else:
    img_data = None

assert batch_size == img_data.shape[0]

Con PyTorch

import glob
import torch
import numpy as np
from PIL import Image
from torchvision import transforms

def _make_3d_tensor(x) -> torch.Tensor:
    """This function is for images that have less channels.

    :param x: input tensor
    :type x: torch.Tensor
    :return: return a tensor with the correct number of channels
    :rtype: torch.Tensor
    """
    return x if x.shape[0] == 3 else x.expand((3, x.shape[1], x.shape[2]))

def preprocess(image, resize_size, crop_size_onnx):
    transform = transforms.Compose([
        transforms.Resize(resize_size),
        transforms.CenterCrop(crop_size_onnx),
        transforms.ToTensor(),
        transforms.Lambda(_make_3d_tensor),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
    
    img_data = transform(image)
    img_data = img_data.numpy()
    img_data = np.expand_dims(img_data, axis=0)
    return img_data

# following code loads only batch_size number of images for demonstrating ONNX inference
# make sure that the data directory has at least batch_size number of images

test_images_path = "automl_models_multi_cls/test_images_dir/*"  # replace with path to images
# Select batch size needed
batch_size = 8
# you can modify resize_size based on your trained model
resize_size = 256
# height and width will be the same for classification
crop_size_onnx = height_onnx_crop_size 

image_files = glob.glob(test_images_path)
img_processed_list = []
for i in range(batch_size):
    img = Image.open(image_files[i])
    img_processed_list.append(preprocess(img, resize_size, crop_size_onnx))
    
if len(img_processed_list) > 1:
    img_data = np.concatenate(img_processed_list)
elif len(img_processed_list) == 1:
    img_data = img_processed_list[0]
else:
    img_data = None

assert batch_size == img_data.shape[0]

Inferencia con ONNX Runtime

La inferencia con ONNX Runtime difiere para cada tarea de Computer Vision.

def get_predictions_from_ONNX(onnx_session, img_data):
    """Perform predictions with ONNX runtime
    
    :param onnx_session: onnx model session
    :type onnx_session: class InferenceSession
    :param img_data: pre-processed numpy image
    :type img_data: ndarray with shape 1xCxHxW
    :return: scores with shapes
            (1, No. of classes in training dataset) 
    :rtype: numpy array
    """

    sess_input = onnx_session.get_inputs()
    sess_output = onnx_session.get_outputs()
    print(f"No. of inputs : {len(sess_input)}, No. of outputs : {len(sess_output)}")    
    # predict with ONNX Runtime
    output_names = [ output.name for output in sess_output]
    scores = onnx_session.run(output_names=output_names,\
                                               input_feed={sess_input[0].name: img_data})
    
    return scores[0]

scores = get_predictions_from_ONNX(session, img_data)

Posprocesamiento

Aplique softmax() en los valores de la predicción para obtener las puntuaciones (probabilidades) de confianza de la clasificación para cada clase. A continuación, la predicción será la clase con la probabilidad más alta.

Sin PyTorch

def softmax(x):
    e_x = np.exp(x - np.max(x, axis=1, keepdims=True))
    return e_x / np.sum(e_x, axis=1, keepdims=True)

conf_scores = softmax(scores)
class_preds = np.argmax(conf_scores, axis=1)
print("predicted classes:", ([(class_idx, classes[class_idx]) for class_idx in class_preds]))

Con PyTorch

conf_scores = torch.nn.functional.softmax(torch.from_numpy(scores), dim=1)
class_preds = torch.argmax(conf_scores, dim=1)
print("predicted classes:", ([(class_idx.item(), classes[class_idx]) for class_idx in class_preds]))

Visualización de predicciones

Visualización de una imagen de entrada con etiquetas.

import matplotlib.image as mpimg
import matplotlib.pyplot as plt
%matplotlib inline

sample_image_index = 0 # change this for an image of interest from image_files list
IMAGE_SIZE = (18, 12)
plt.figure(figsize=IMAGE_SIZE)
img_np = mpimg.imread(image_files[sample_image_index])

img = Image.fromarray(img_np.astype('uint8'), 'RGB')
x, y = img.size

fig,ax = plt.subplots(1, figsize=(15, 15))
# Display the image
ax.imshow(img_np)

label = class_preds[sample_image_index]
if torch.is_tensor(label):
    label = label.item()
    
conf_score = conf_scores[sample_image_index]
if torch.is_tensor(conf_score):
    conf_score = np.max(conf_score.tolist())
else:
    conf_score = np.max(conf_score)

display_text = '{} ({})'.format(label, round(conf_score, 3))
print(display_text)

color = 'red'
plt.text(30, 30, display_text, color=color, fontsize=30)

plt.show()

Pasos siguientes