pytorch深度学习入门（12）之-神经网络导出onnx模型部署与应用

码农呆呆

已于 2023-11-30 11:10:35 修改

阅读量1k

点赞数 14

分类专栏：深度学习文章标签：深度学习 pytorch 神经网络

于 2023-11-15 17:58:07 首次发布

本文链接：https://blog.csdn.net/ajunbin859/article/details/134425726

版权

本文介绍了如何将PyTorch训练的深度学习模型导出为ONNX格式，讨论了ONNX的优势和导出过程。通过使用ONNX Runtime进行推理，验证了模型的输出一致性。文章涵盖了依赖安装、模型导出、ONNX模型验证和在实际图像上的应用。

摘要由CSDN通过智能技术生成

概述：

ONNX（Open Neural Network Exchange）是一种开放神经网络交换格式，它使得不同深度学习框架（如TensorFlow、PyTorch、MXNet等）之间的互操作成为可能。ONNX提供了一种标准化的方式，可以将训练好的模型导出并转换为ONNX格式，然后可以在其他支持ONNX的框架或工具中进行部署和推理。

ONNX的主要优势在于它促进了深度学习模型在不同平台之间的互操作性和可移植性。这对于跨平台开发和部署深度学习应用非常有用，因为它可以减少在不同框架之间迁移模型所需的工作量，并帮助加速深度学习技术的普及和应用。

要了解更多关于ONNX的信息，可以访问ONNX的官方网站。

开放神经网络 eXchange (ONNX)是一种用于表示机器学习模型的开放标准格式。该torch.onnx模块提供 API 来从本机 PyTorch 模型捕获计算图torch.nn.Module并将其转换为ONNX 图。

导出的模型可由支持 ONNX 的众多运行时中的任何一个使用，包括 Microsoft 的ONNX Runtime。

注意：

=目前，ONNX 导出器 API 有两种风格，但本教程将重点介绍torch.onnx.dynamo_export.=

TorchDynamo 引擎用于连接 Python 的框架评估 API 并将其字节码动态重写为FX 图表。生成的 FX 图表在最终转换为 ONNX 图表之前经过打磨。

这种方法的主要优点是使用字节码分析捕获FX 图，从而保留模型的动态特性，而不是使用传统的静态跟踪技术。

依赖关系

需要 PyTorch 2.1.0 或更高版本。

ONNX 导出器依赖于额外的 Python 包：

ONNX标准库

ONNX 脚本库使开发人员能够使用 Python 子集以富有表现力且简单的方式编写 ONNX 运算符、函数和模型。

它们可以通过pip安装：

pip install --upgrade onnx onnxscript

要验证安装，请运行以下命令：

import torch
print(torch.__version__)

import onnxscript
print(onnxscript.__version__)

from onnxscript import opset18  # opset 18 is the latest (and only) supported version for now

import onnxruntime
print(onnxruntime.__version__)

每次导入都必须成功且没有任何错误，并且必须打印出库版本。

接下来讲述如何运行onnx模型
1.将 PyTorch 模型导出到 ONNX
2.扩展ONNX注册表

导出的模型将使用 ONNX Runtime 执行。ONNX Runtime 是 ONNX 模型的注重性能的引擎，可跨多个平台和硬件（Windows、Linux 和 Mac，以及 CPU 和 GPU）进行高效推理。事实证明，ONNX Runtime 可以显着提高多个模型的性能，如此处所述

对于本教程，您需要安装ONNX 和ONNX Runtime。您可以使用以下命令获取 ONNX 和 ONNX Runtime 的二进制版本
安装onnx运行环境：

pip install onnx onnxruntime

ONNX 运行时建议使用 PyTorch 的最新稳定运行时。

# Some standard imports
import numpy as np

from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.onnx

超分辨率是一种提高图像、视频分辨率的方法，广泛应用于图像处理或视频编辑中。在本教程中，我们将使用一个小型超分辨率模型。

首先，让我们SuperResolution在 PyTorch 中创建一个模型。该模型使用 Shi 等人在“使用高效子像素卷积神经网络的实时单图像和视频超分辨率”中描述的高效子像素卷积层，以提高图像的分辨率。该模型期望图像的 Y 分量YCbCr作为输入，并以超分辨率输出放大的 Y 分量。

该模型直接来自 PyTorch 的示例，未经修改：

# Super Resolution model definition in PyTorch
import torch.nn as nn
import torch.nn.init as init


class SuperResolutionNet(nn.Module):
    def __init__(self, upscale_factor, inplace=False):
        super(SuperResolutionNet, self).__init__()

        self.relu = nn.ReLU(inplace=inplace)
        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2)