使用 ONNX 对来自 AutoML 的计算机视觉模型进行预测

适用于:Python SDK azure-ai-ml v2(当前版本)

本文将介绍如何使用 Open Neural Network Exchange (ONNX) 对从 Azure 机器学习中的自动机器学习 (AutoML) 生成的计算机视觉模型进行预测。

若要使用 ONNX 进行预测,需执行以下操作:

  1. 从 AutoML 训练运行下载 ONNX 模型文件。
  2. 了解 ONNX 模型的输入和输出。
  3. 预处理数据,使其成为输入图像所需的格式。
  4. 对 Python 的 ONNX 运行时执行推理。
  5. 可视化对象检测和实例分段任务的预测。

ONNX 是机器学习和深度学习模型的开放标准。 它支持在常用的 AI 框架中进行模型导入和导出(互操作性)。 更多详细信息,请浏览 ONNX GitHub 项目

ONNX 运行时是一个支持跨平台推理的开源项目。 ONNX 运行时提供跨编程语言(包括 Python、C++、C#、C、Java 和 JavaScript)的 API。 可以使用这些 API 对输入图像执行推理。 将模型导出为 ONNX 格式后,便可以在项目需要的任何编程语言中使用这些 API。

本指南介绍如何使用用于 ONNX 运行时的 Python API 对常见视觉任务的图像进行预测。 可以跨语言使用这些 ONNX 导出的模型。

先决条件

下载 ONNX 模型文件

可以使用 Azure 机器学习工作室 UI 或 Azure 机器学习 Python SDK 从 AutoML 运行下载 ONNX 模型文件。 建议使用具有实验名称和父运行 ID 的 SDK 进行下载。

Azure 机器学习工作室

在 Azure 机器学习工作室中,通过训练笔记本中生成的指向实验的超链接进入实验,或选择“资产”下的“实验”选项卡中实验名称进入实验 。 然后,选择最佳子运行。

在最佳子运行中,转到“输出+日志”>“train_artifacts” 。 使用“下载”按钮手动下载以下文件:

  • labels.json:包含训练数据集中所有类或标签的文件。
  • model.onnx:ONNX 格式的模型。

显示下载 ONNX 模型文件的选项的屏幕截图。

将下载的模型文件保存到目录。 本文中的示例使用 ./automl_models 目录。

Azure 机器学习 Python SDK

在 SDK 中,可以使用实验名称和父运行 ID 选择最佳子运行(按主要指标)。 然后,可以下载 labels.json 和 model.onnx 文件 。

以下代码根据相关的主要指标返回最佳子运行。

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)

下载 labels.json 文件,其中包含训练数据集中的所有类和标签。

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
)

下载 model.onnx 文件。

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

如果使用 ONNX 模型对对象检测和实例分段进行批量推理,请参阅有关批量评分的模型生成的部分。

生成模型进行批量评分

默认情况下,AutoML for Images 支持分类的批量评分。 但是对象检测和实例分割 ONNX 模型不支持批量推理。 若要对于对象检测和实例分段进行批量推断,请使用以下过程为所需的批大小生成 ONNX 模型。 为特定批大小生成的模型不能用于其他批大小。

下载 conda 环境文件,并创建用于命令作业的环境对象。

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

使用以下模型特定参数提交脚本。 有关参数的更多详细信息,请参阅模型特定超参数;有关受支持的对象检测模型名称,请参阅受支持的模型体系结构部分

若要获取创建批处理评分模型所需的参数值,请参阅 AutoML 训练运行 outputs 文件夹下生成的评分脚本。 使用最佳子运行评分文件内模型设置变量中提供的超参数值。

对于多类图像分类,为最佳子运行生成的 ONNX 模型默认支持批量评分。 因此,此任务类型不需要模型特定的参数,可以跳到加载标签和 ONNX 模型文件部分。

下载 ONNX_batch_model_generator_automl_for_images.py 文件并将其保留在当前目录中,然后提交脚本。 使用以下命令作业提交 azureml-examples GitHub 存储库中可用的脚本 ONNX_batch_model_generator_automl_for_images.py,以生成特定批大小的 ONNX 模型。 在以下代码中,训练的模型环境用来提交此脚本,以生成 ONNX 模型并将其保存到 outputs 目录。

对于多类图像分类,为最佳子运行生成的 ONNX 模型默认支持批量评分。 因此,此任务类型不需要模型特定的参数,可以跳到加载标签和 ONNX 模型文件部分。

生成批模型后,可手动通过 UI 从 Outputs+>logsoutputs 中下载,或者使用以下方法:

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
)

完成模型下载步骤后,使用 ONNX 运行时 Python 包并使用 model.onnx 文件执行推理。 为进行演示,本文使用如何为每个视觉任务准备图像数据集中的数据集。

我们已使用所有视觉任务各自的数据集为其训练了模型,以演示 ONNX 模型推理。

加载标签和 ONNX 模型文件

以下代码片段加载 labels.json,其中类名已排序。 也就是说,如果 ONNX 模型预测标签 ID 为 2,则它对应于 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))

获取 ONNX 模型的预期输入和输出详细信息

使用模型时,务必了解一些特定于模型和特定于任务的详细信息。 这些详细信息包括输入数量和输出数量、用于预处理图像的预期输入形状或格式,以及输出形状,确保你了解特定于模型或特定于任务的输出。

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}") 

ONNX 模型的预期输入和输出格式

每个 ONNX 模型都有一组预定义的输入和输出格式。

此示例应用具有 134 个图像和 4 个类/标签的 fridgeObjects 数据集上训练的模型,以说明 ONNX 模型推理。 有关训练图像分类任务的详细信息,请参阅多类图像分类笔记本

输入格式

输入是经过预处理的图像。

输入名称 输入形状 输入类型 描述
input1 (batch_size, num_channels, height, width) ndarray(float) 输入是经过预处理的图像,形状为 (1, 3, 224, 224),批大小为 1,高度和宽度为 224。 这些数字对应于训练示例中 crop_size 所用的值。

输出格式

输出是所有类/标签的 logit 数组。

输出名称 输出形状 输出类型 描述
output1 (batch_size, num_classes) ndarray(float) 模型返回 logit(没有 softmax)。 例如,对于批大小为 1 和 4 的类,它返回 (1, 4)

预处理

执行以下预处理步骤,以实现 ONNX 模型推理:

  1. 将图像转换为 RGB。
  2. 将图像大小调整为 valid_resize_sizevalid_resize_size 值,这些值对应于训练期间验证数据集转换时使用的值。 valid_resize_size 的默认值为 256。
  3. 将图像中心裁剪为 height_onnx_crop_sizewidth_onnx_crop_size。 它与 valid_crop_size 对应,默认值为 224。
  4. HxWxC 更改为 CxHxW
  5. 转换为 float 型。
  6. 使用 ImageNet 的 mean = [0.485, 0.456, 0.406]std = [0.229, 0.224, 0.225] 进行规范化。

如果在训练期间为超参数 valid_resize_sizevalid_crop_size 选择了不同的值,则应使用这些值。

获取 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

无 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]

有 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]

使用 ONNX 运行时进行推理

使用 ONNX 运行时进行推理因各个计算机视觉任务而异。

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)

后期处理

softmax() 应用预测值,以获取每个类的分类置信度分数(概率)。 然后,将预测出概率最高的类。

无 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]))

有 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]))

将预测结果可视化

使用标签将输入图像可视化。

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

后续步骤