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

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

本文介绍如何使用开放神经网络Exchange(ONNX)对Azure 机器学习中的自动化机器学习(AutoML)生成的计算机视觉模型进行预测。

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

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

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

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

本指南介绍如何使用 Python API for ONNX Runtime对图像进行预测,以便执行热门视觉任务。 可以跨语言使用这些 ONNX 导出的模型。

先决条件

下载 ONNX 模型文件

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

Azure 机器学习工作室

在Azure 机器学习工作室上,通过使用训练笔记本中生成的试验的超链接转到试验,或在 Experiments 选项卡上选择Assets下的实验名称。 然后,选择最佳子运行。

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

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

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

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

Azure 机器学习 Python SDK

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

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

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

try:
    credential = DefaultAzureCredential()
    credential.get_token("https://management.chinacloudapi.cn/.default")
except Exception:
    credential = InteractiveBrowserCredential()
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 文件并将其保留在当前目录中,然后提交脚本。 使用以下命令提交ONNX_batch_model_generator_automl_for_images.py脚本,该脚本可在azureml-examples GitHub 仓库中获得,以生成指定批大小的 ONNX 模型。 在以下代码中,训练的模型环境用来提交此脚本,以生成 ONNX 模型并将其保存到 outputs 目录。

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

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

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

后续步骤