PyTorch模子转换ONNX 入门

目次

前言
什么是ONNX文件
ONNX文件简介
ONNX 文件的重要特点
ONNX 文件的根本布局
pytorch模子转ONNX模子
前期预备
cuda安装
pytorch安装
ONNX模块安装
pytorch模子导出ONNX模子 
简朴体验
检察ONNX文件
 torch.onnx.export函数
 加载ONNX模子

---

前言

什么是ONNX文件

ONNX文件简介

ONNX（Open Neural Network Exchange）是一种开放的文件格式，用于表现呆板学习模子，旨在促进差异框架之间的互操纵性。
ONNX 文件通常以 .onnx 为扩展名，可以大概存储神经网络的布局和权重，使得模子可以在差异的深度学习框架（如 TensorFlow、PyTorch、Caffe 等）之间举行转换和摆设。
ONNX 文件的重要特点

• 跨平台兼容性：ONNX 支持在多个框架之间共享模子，用户可以在一个框架中训练模子，然后将其导出为 ONNX 格式，以便在另一个框架中举行推理。
  
• 开放标准：ONNX 是一个开放的标准，由多个行业相助同伴共同开发和维护。它为呆板学习社区提供了一个同一的格式。
  
• 高效性：ONNX 文件可以大概有效地存储模子的盘算图、参数和操纵，如许可以更高效地举行推理。
  
• 支持多种操纵：ONNX 界说了一组标准操纵符，支持多种神经网络架构，包罗卷积神经网络（CNN）、循环神经网络（RNN）等。
  
• 工具支持：ONNX 提供了一系列工具和库，支持将模子从差异框架导出到 ONNX 格式，也支持从 ONNX 文件加载模子举行推理。
  

ONNX 文件的根本布局

一个范例的 ONNX 文件包罗以下内容：

• 盘算图：形貌了模子的布局，包罗各层的毗连关系。
  
• 参数：模子的权重和偏置值。
  
• 元数据：关于模子的一些附加信息，如输入输出的外形、数据范例等。
  

pytorch模子转ONNX模子

前期预备

用户需事先安装cuda、cudnn（可选）和pytorch
cuda安装

windows下cuda的安装见
windows安装cuda与cudnn-CSDN博客
linux下的cuda安装见
【CUDA】Ubuntu体系怎样安装CUDA保姆级教程(2022年最新)_ubuntu安装cuda-CSDN博客
无论在哪个体系上安装cuda，只要输入以下下令时有信息输出即表现安装乐成

1. nvcc -V

复制代码

pytorch安装

当cuda安装乐成后，输入nvcc -V下令检察cuda版本号，然后进入pytorch官网，下载对应cuda版本的pytorch即可
Previous PyTorch Versions | PyTorch
无论什么体系，只要在下令行输出以下结果，即表现pytorch安装乐成

ONNX模块安装

1. pip install onnx

复制代码

1. pip install onnxruntime

复制代码

pytorch模子导出ONNX模子 

简朴体验

起首利用pytorch写一个简朴的网络模子

1. import torch
2. import torchvision
3. import numpy as np
5. devide=torch.device("cuda" if torch.cuda.is_available() else "cpu")
6. # 定义一个简单的PyTorch 模型
7. class MyModel(torch.nn.Module):
8.     def __init__(self):
9.         super(MyModel, self).__init__()
10.         self.conv1 = torch.nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
11.         self.relu = torch.nn.ReLU()
12.         self.maxpool = torch.nn.MaxPool2d(kernel_size=2, stride=2)
13.         self.conv2 = torch.nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
14.         self.flatten = torch.nn.Flatten()
15.         self.fc1 = torch.nn.Linear(64 * 8 * 8, 10)
17.     def forward(self, x):
18.         x = self.conv1(x)
19.         x = self.relu(x)
20.         x = self.maxpool(x)
21.         x = self.conv2(x)
22.         x = self.relu(x)
23.         x = self.maxpool(x)
24.         x = self.flatten(x)
25.         x = self.fc1(x)
26.         return x
28. # 创建模型实例
29. model = MyModel().to(devide)
31. # 指定模型输入尺寸
32. dummy_input = torch.randn(1, 3, 32, 32).to(devide)
34. # 将PyTorch模型转为ONNX模型
35. torch.onnx.export(model, dummy_input, 'mymodel.onnx',  do_constant_folding=False)

复制代码

如上所示，我们手写了一个简朴的网络模子，可以看到，转ONNX模子文件的代码只有末了一行

1. torch.onnx.export(model, dummy_input, 'mymodel.onnx',  do_constant_folding=False)

复制代码

毕竟上，torch.onnx.export函数就是将一个torch模子转换为ONNX文件的函数
检察ONNX文件

上述代码乐成运行后，会在当地天生一个mymodel.onnx文件，该文件的打开须要利用netron，有关netron的安装见
netron安装（windows && linux）-CSDN博客
安装乐成后，利用netron打开mymodel.onnx，如下所示

 torch.onnx.export函数

上述我们写了一个小demo体验了torch模子转换为ONNX文件，并检察了ONNX文件到底是什么，接下来我们来看torch模子转换ONNX文件的核心函数中参数寄义都是什么

1. torch.onnx.export(
2.             model,
3.             args,
4.             f,
5.             export_params=True,
6.             opset_version=10,
7.             do_constant_folding=True,
8.             input_names=['input'],
9.             output_names=['output'],
10.             dynamic_axes=None,
11.             verbose=False,
12.             example_outputs=None,
13.             keep_initializers_as_inputs=None)

复制代码

  参数详解
  

• model:
  
  
  
  • 范例: torch.nn.Module
    
  • 形貌: 被转换的 PyTorch 模子。
    
  
  
• args:
  
  
  
  • 范例: tuple 或 torch.Tensor
    
  • 形貌: torch模子的输入示例，可以是一个单一的张量或多个张量（以元组的情势）。这些输入数据用于实验模子，确定模子的输入外形。
    
  
  
• f:
  
  
  
  • 范例: str 或 Path
    
  • 形貌: 导出模子的目的文件路径或文件名，通常以 .onnx 作为扩展名。
    
  
  
• export_params:
  
  
  
  • 范例: bool，默认: True
    
  • 形貌: 是否将模子的参数（权重和偏置）也导出到 ONNX 文件中。假如设置为 True，导出的模子会包罗全部的参数。
    
  
  
• opset_version:
  
  
  
  • 范例: int，默认: 9
    
  • 形貌: 指定要利用的 ONNX 操纵集版本。差异版本大概支持差异的操纵和功能。设置符合的版本可以确保兼容性。
    
  
  
• do_constant_folding:
  
  
  
  • 范例: bool，默认: True
    
  • 形貌: 是否举行常量折叠优化。常量折叠会在导出过程中将一些常量盘算提前，从而简化模子的盘算图，提升推理服从。
    
  
  
• input_names:
  
  
  
  • 范例: list，默认: None
    
  • 形貌: 输入张量的名称列表。可以用来指定导出模子输入的名称，便于后续在其他框架中辨认。
    
  
  
• output_names:
  
  
  
  • 范例: list，默认: None
    
  • 形貌: 输出张量的名称列表。类似于 input_names，用于指定导出模子输出的名称。
    
  
  
• dynamic_axes:
  
  
  
  • 范例: dict 或 None，默认: None
    
  • 形貌: 允许动态维度的输入输出。在导出时，可以指定某些维度是动态的，如许在推理时输入的外形可以变革。比方：
    1. dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}

    复制代码
    这表现 input 和 output 的第一个维度是动态的（比方 batch size）。
    
  
  
• verbose:
  
  
  
  • 范例: bool，默认: False
    
  • 形貌: 是否在导出时打印详细信息。假如设置为 True，会表现更多的调试信息，便于跟踪导出过程中的标题。
    
  
  
• example_outputs:
  
  
  
  • 范例: tuple 或 torch.Tensor，默认: None
    
  • 形貌: 用于指定模子的示例输出，这有助于 ONNX 在导出时举行范例推断。可以提供一个或多个输出张量，以便更好地推断输出的外形和范例。
    
  
  
• keep_initializers_as_inputs:
  
  
  
  • 范例: bool，默认: None
    
  • 形貌: 是否将模子的初始值（权重）作为输入生存。假如设置为 True，则初始值将被视为模子的输入之一，而不是存储在模子的参数中。
    
  
  

下面是一个常用的模板

1. import torch.onnx
3. # 转为ONNX
4. def Convert_ONNX(model):
6.     # 设置模型为推理模式
7.     model.eval()
9.     # 设置模型输入的尺寸
10.     dummy_input = torch.randn(1, input_size, requires_grad=True)  
12.     # 导出ONNX模型  
13.     torch.onnx.export(model,         # model being run
14.          dummy_input,       # model input (or a tuple for multiple inputs)
15.          "xxx.onnx",       # where to save the model  
16.          export_params=True,  # store the trained parameter weights inside the model file
17.          opset_version=10,    # the ONNX version to export the model to
18.          do_constant_folding=True,  # whether to execute constant folding for optimization
19.          input_names = ['modelInput'],   # the model's input names
20.          output_names = ['modelOutput'], # the model's output names
21.          dynamic_axes={'modelInput' : {0 : 'batch_size'},    # variable length axes
22.                                 'modelOutput' : {0 : 'batch_size'}})
23.     print(" ")
24.     print('Model has been converted to ONNX')
27. if __name__ == "__main__":
29.     # 构建模型并训练
30.     # xxxxxxxxxxxx
32.     # 测试模型精度
33.     #testAccuracy()
35.     # 加载模型结构与权重
36.     model = Network()
37.     path = "myFirstModel.pth"
38.     model.load_state_dict(torch.load(path))
40.     # 转换为ONNX
41.     Convert_ONNX(model)

复制代码

 加载ONNX模子

导出ONNX模子后，加载ONNX模子须要用到onnxruntime库，以下是一个导出ONNX模子的示例

1. import onnxruntime as ort
3. # 加载 ONNX 模型
4. ort_session = ort.InferenceSession("model.onnx")
6. # 准备输入信息
7. input_info = ort_session.get_inputs()[0]
8. input_name = input_info.name
9. input_shape = input_info.shape
10. input_type = input_info.type
13. # 运行ONNX模型
14. outputs = ort_session.run(input_name, input_data)
16. # 获取输出信息
17. output_info = ort_session.get_outputs()[0]
18. output_name = output_info.name
19. output_shape = output_info.shape
20. output_data = outputs[0]
22. print("outputs:", outputs)
23. print("output_info :", output_info )
24. print("output_name :", output_name )
25. print("output_shape :", output_shape )
26. print("output_data :", output_data )

复制代码

在以下案例中，我们起首将resnet-18模子导出为ONNX模子，然后再加载导出的ONNX模子，末了对比torch模子和ONNX模子的输出差异

1. import torch
2. import torchvision.models as models
3. import onnx
4. import onnxruntime
6. # 加载 PyTorch 模型
7. model = models.resnet18(pretrained=True)
8. model.eval()
10. # 定义输入和输出张量的名称和形状
11. input_names = ["input"]
12. output_names = ["output"]
13. batch_size = 1
14. input_shape = (batch_size, 3, 224, 224)
15. output_shape = (batch_size, 1000)
17. # 将 PyTorch 模型转换为 ONNX 格式
18. torch.onnx.export(
19.     model,  # 要转换的 PyTorch 模型
20.     torch.randn(input_shape),  # 模型输入的随机张量
21.     "resnet18.onnx",  # 保存的 ONNX 模型的文件名
22.     input_names=input_names,  # 输入张量的名称
23.     output_names=output_names,  # 输出张量的名称
24.     dynamic_axes={input_names[0]: {0: "batch_size"}, output_names[0]: {0: "batch_size"}}  # 动态轴，即输入和输出张量可以具有不同的批次大小
25. )
27. # 加载 ONNX 模型
28. onnx_model = onnx.load("resnet18.onnx")
29. onnx_model_graph = onnx_model.graph
30. onnx_session = onnxruntime.InferenceSession(onnx_model.SerializeToString())
32. # 使用随机张量测试 ONNX 模型
33. x = torch.randn(input_shape).numpy()
34. onnx_output = onnx_session.run(output_names, {input_names[0]: x})[0]
36. print(f"PyTorch output: {model(torch.from_numpy(x)).detach().numpy()[0, :5]}")
37. print(f"ONNX output: {onnx_output[0, :5]}")

复制代码

运行结果如下所示

1. PyTorch output: [0.22972351 2.4930785  2.4462368  2.7443404  4.7080407 ]
2. ONNX output: [0.22972152 2.4930775  2.4462373  2.7443395  4.708042  ]

复制代码