ONNX Runtime 部署语义分割模型

使用 ONNX Runtime 部署语义分割模型

在深度学习应用中，语义分割是图像处理中非常重要的一项任务。语义分割的目标是将图像中的每个像素分类到一个特定的类别中。随着深度学习框架的发展，ONNX（Open Neural Network Exchange）作为一种开源格式，逐渐成为了跨平台深度学习模型的标准化表达。ONNX Runtime 是微软提供的高效推理引擎，支持在多种硬件平台上高效执行ONNX格式的模型。本文将介绍如何使用 ONNX Runtime 部署一个语义分割模型。

1. 什么是 ONNX Runtime？

ONNX Runtime 是一个跨平台的高性能推理引擎，支持多个硬件平台，如 CPU、GPU、以及一些定制硬件。ONNX 作为开放的深度学习框架标准，可以将来自不同深度学习框架（如 PyTorch、TensorFlow、Keras 等）的模型转换为 ONNX 格式。ONNX Runtime 的优势在于能够为多种平台提供高效的推理执行，并且可以轻松集成到生产环境中。

2. 语义分割模型概述

语义分割是计算机视觉中的一项重要任务，目标是为图像中的每个像素分配一个标签，通常应用于医学影像分析、自动驾驶、卫星图像分析等领域。深度学习模型，如 U-Net 和 DeepLabV3，已经在该任务中取得了巨大的成功。

3. 使用 ONNX Runtime 进行模型推理

步骤 1：准备语义分割模型

假设已经训练好一个语义分割模型（例如，基于 U-Net 或 DeepLabV3），并将其转换为 ONNX 格式。如果你使用的是 PyTorch 或 TensorFlow，可以将训练好的模型转成onnx格式，然后用下面推理脚本进行推理。

步骤 2：安装 ONNX Runtime

首先，需要在系统中安装 ONNX Runtime。可以使用以下命令进行安装：

pip install onnxruntime

如果你想启用 GPU 加速，可以安装 GPU 版本的 ONNX Runtime：

pip install onnxruntime-gpu

步骤 3：加载和运行 ONNX 模型

在 ONNX Runtime 中加载并推理模型非常简单。已经将模型转换为 end2end.onnx 格式，接下来，可以使用以下代码来进行推理。

import cv2
import numpy as np
import matplotlib.pyplot as plt
import onnxruntime as ort

# 参数
IMG_FPATH = "./test/1.png"
ONNX_MODEL_FPATH = "./model/end2end.onnx"
num_classes = 1  # 类别数

# 加载图像
img = cv2.imread(IMG_FPATH)
if img is None:
    raise Exception('Image not found!')
# RGB2BGR
rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# 输出img的尺寸
print(rgb_img.shape)

# 加载ONNX模型
session = ort.InferenceSession(ONNX_MODEL_FPATH)

# 获取输入和输出名称
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 归一化方法
mean = np.array([0.20, 0.19, 0.17])
std = np.array([0.15, 0.14, 0.14])

# 模型的输入数据准备
input_shape = session.get_inputs()[0].shape

# 调整图像尺寸以匹配输入形状
resized_img = cv2.resize(img, (input_shape[3], input_shape[2]))
# 输出resized_img的尺寸
# print(resized_img.shape)

# np.transpose(resized_img, (2, 0, 1))：对 resized_img 数组的维度进行转置操作。OpenCV 图像的原始维度顺序是（高度、宽度、通道），但是神经网络通常希望在高度和宽度之前包含通道维度。因此，这一步将维度重新排列为（通道、高度、宽度）
# .astype(np.float32) / 255.0：这一步将转置和调整大小后的图像的像素值从整数（通常在 0 到 255 范围内）转换为浮点数，范围从 0.0 到 1.0
preprocessed_img = np.transpose(resized_img, (2, 0, 1)).astype(np.float32) / 255.0  # Normalize to [0, 1]

# 标准化
preprocessed_img[0] = (preprocessed_img[0] - mean[0]) / std[0]
preprocessed_img[1] = (preprocessed_img[1] - mean[1]) / std[1]
preprocessed_img[2] = (preprocessed_img[2] - mean[2]) / std[2]

# 扩展维度以匹配输入形状 
preprocessed_img = np.expand_dims(preprocessed_img, axis=0)

# 执行推理
outputs = session.run([output_name], {input_name: preprocessed_img})
output_data = outputs[0].reshape((input_shape[2], input_shape[3]))

# 获取每一个类别的像素值
classIds = output_data.astype(np.uint8)
# print(classIds)

# 自定义mask
COLORS = np.zeros((255, 3), dtype="uint8")
# for i in range(num_classes):
#     COLORS[i] = np.random.randint(0, 256, size=3)
# Background
COLORS[0] = [0, 0, 0]
# Class 1
COLORS[1] = [0, 200, 200]
# # Class 2
# COLORS[2] = [0, 0, 200]

# 处理输出数据
segment_map = COLORS[classIds]
segment_map = cv2.resize(segment_map, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_NEAREST)

# Mix original with segmentation map
img_with_segmentation = (0.4 * segment_map + 0.6 * img).astype(np.uint8)

# 可视化结果
def show_img(img):
    dpi = 80
    height, width, _ = img.shape
    figsize = width / float(dpi), height / float(dpi)

    fig = plt.figure(figsize=figsize)
    plt.axis('off')
    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    plt.show()

show_img(img_with_segmentation)

步骤 4：优化和部署

ONNX Runtime 支持多种优化方法，可以提高推理速度，特别是在生产环境中，优化非常重要。以下是一些优化选项：

1. 使用 ONNX Runtime 的图优化：

ONNX Runtime 提供了自动化的图优化功能，通过 optimize 选项来减少计算量，提高效率。可以通过以下代码启用图优化：

session = ort.InferenceSession("end2end.onnx", providers=['CPUExecutionProvider'])

2. 使用硬件加速：

ONNX Runtime 支持多种硬件加速，如 NVIDIA GPU 或 Intel OpenVINO。可以根据实际硬件进行配置，以下是启用 GPU 加速的示例：

session = ort.InferenceSession("end2end.onnx", providers=['CUDAExecutionProvider'])

步骤 5：批量推理

对于需要处理大量图像的场景，可以使用 ONNX Runtime 批量推理的功能，以提高性能。在执行推理时，可以通过一次传入多个图像来进行批量处理。

batch_images = np.stack([image1, image2, image3], axis=0)  # 假设有三个输入图像
outputs = session.run(None, {input_name: batch_images})

4. 总结

ONNX Runtime 是一个高效、跨平台的推理引擎，可以帮助开发者在各种硬件平台上运行 ONNX 格式的深度学习模型。通过将语义分割模型转换为 ONNX 格式并使用 ONNX Runtime 进行推理，可以获得更好的性能，尤其是在生产环境中。本文介绍了如何将训练好的语义分割模型转换为 ONNX 格式、如何使用 ONNX Runtime 进行推理以及如何优化模型性能。随着深度学习技术的不断发展，ONNX 和 ONNX Runtime 提供了一个强大的工具集，帮助开发者将深度学习模型高效部署到生产环境中。