PyTorch学习（6）：导出ONNX模型

1. 简述

        在使用PyTorch进行模型训练时，我们通常希望将训练好的模型导出为ONNX（Open Neural Network Exchange）格式的文件。

        ONNX（Open Neural Network Exchange）是一个开放格式，用于表示深度学习模型。它允许模型被不同的深度学习框架所使用，实现了模型的互操作性和可移植性。ONNX由Facebook和微软在2016年共同推出，旨在解决深度学习模型在不同框架和工具之间的兼容性问题。

2. 工作原理及用途

        ONNX定义了一种模型表示方法，包括一系列的操作（operators）和数据结构。模型被表示为一个计算图，图中的节点是操作，边是数据流。ONNX模型包含两部分：

模型定义：描述了模型的结构，包括层、节点和连接方式。

模型权重：包含了模型训练过程中学习到的参数，如卷积核、偏置项等。

ONNX使用JSON格式来描述模型，这使得模型易于理解和修改。同时，ONNX还提供了一系列的工具，如ONNX Checker、ONNX Runtime等，用于模型的验证和推理。

ONNX有哪些用途呢？

模型共享：ONNX格式使得模型可以在不同的深度学习社区和项目之间共享，从而促进了知识的传播和协作。

模型转换：ONNX提供了一个中间格式，可以将一个框架训练的模型转换为ONNX格式，然后在其他支持ONNX的框架中使用。

跨平台部署：ONNX支持多种硬件和软件平台，使得模型可以轻松地部署到不同的设备和操作系统上，包括服务器、移动设备和嵌入式系统。

模型优化：ONNX模型可以被优化和精简，以适应特定的部署环境和性能要求。

生态系统支持：ONNX生态系统提供了一系列的工具和库，支持模型的验证、推理、可视化和转换。

多样化的应用场景：ONNX模型可以应用于各种领域，如计算机视觉、自然语言处理、语音识别等。

3. PyTorch导出ONNX模型文件

        在导出模型文件之前，需要确认的是，导出的是经过训练的模型（未经训练导出的模型没有什么意义），而且需要导出的是计算图和权重参数模型，不能导出state_dict。

准备数据

# 创建一个示例输入张量

dummy_input = torch.randn(1, 3, 224, 224)  # 假设输入图像的大小为224x224，3个通道

        当前准备的数据是一个模拟数据，并不是要真实使用的数据，因此可以采用随机生成的形式获得。

导出ONNX模型

# 导出模型

torch.onnx.export(model,               # 要导出的模型

                  dummy_input,         # 模型的输入数据

                  "model.onnx",         # 要保存的文件名

                  export_params=True,   # 是否导出模型参数

                  opset_version=11,     # ONNX版本

                  do_constant_folding=True,  # 是否执行常量折叠优化

                  input_names=['input'],   # 输入数据的名字

                  output_names=['output'], # 输出数据的名字

                  dynamic_axes={'input': {0: 'batch_size'},  # 输入数据的动态轴

                                'output': {0: 'batch_size'}})

        在使用 torch.onnx.export 函数导出 PyTorch 模型为 ONNX 格式时，input_names 参数用于为模型的输入张量指定名称。这些名称在 ONNX 模型中用于标识输入数据的节点，并且在使用 ONNX 模型进行推理时，这些名称会被用作输入数据的键。

        input_names 应该是一个列表，列表中的每个元素都是一个字符串，对应于模型输入张量的名称。