PyTorch模型转换为ONNX格式实现过程浅谈

浅浅的花香味﹌ 2024-03-24 18:47 79阅读 0赞

![d1f5eadc79bdaef445efcee8564f8fbb.jpeg][]

本文主要为大家介绍了PyTorch模型转换为ONNX格式实现过程详解，有需要的朋友可以借鉴参考下，希望能够有所帮助！

### 1. 安装依赖 ###

将PyTorch模型转换为ONNX格式可以使它在其他框架中使用，如TensorFlow、Caffe2和MXNet

首先安装以下必要组件：

*  Pytorch
 *  ONNX
 *  ONNX Runtime（可选）

建议使用`conda`环境，运行以下命令来创建一个新的环境并激活它：

<table> 
 <tbody> 
  <tr> 
   <td> <p>1</p> <p>2</p> </td> 
   <td> <p><code>conda create </code><code>-</code><code>n onnx python</code><code>=</code><code>3.8</code></p> <p><code>conda activate onnx</code></p> </td> 
  </tr> 
 </tbody> 
</table>

接下来使用以下命令安装PyTorch和ONNX：

<table> 
 <tbody> 
  <tr> 
   <td> <p>1</p> <p>2</p> </td> 
   <td> <p><code>conda install pytorch torchvision torchaudio </code><code>-</code><code>c pytorch</code></p> <p><code>pip install onnx</code></p> </td> 
  </tr> 
 </tbody> 
</table>

可选地，可以安装ONNX Runtime以验证转换工作的正确性：

<table> 
 <tbody> 
  <tr> 
   <td> <p>1</p> </td> 
   <td> <p><code>pip install onnxruntime</code></p> </td> 
  </tr> 
 </tbody> 
</table>

### 2. 准备模型 ###

将需要转换的模型导出为PyTorch模型的`.pth`文件。使用PyTorch内置的函数加载它，然后调用eval()方法以保证close状态：

<table> 
 <tbody> 
  <tr> 
   <td> <p>1</p> <p>2</p> <p>3</p> <p>4</p> <p>5</p> <p>6</p> <p>7</p> <p>8</p> <p>9</p> <p>10</p> <p>11</p> <p>12</p> <p>13</p> <p>14</p> <p>15</p> <p>16</p> <p>17</p> <p>18</p> <p>19</p> <p>20</p> <p>21</p> <p>22</p> <p>23</p> <p>24</p> <p>25</p> <p>26</p> <p>27</p> <p>28</p> <p>29</p> </td> 
   <td> <p><code>import</code> <code>torch.nn as nn</code></p> <p><code>import</code> <code>torch.nn.functional as F</code></p> <p><code>import</code> <code>torch.optim as optim</code></p> <p><code>import</code> <code>torch.onnx</code></p> <p><code>import</code> <code>torchvision.transforms as transforms</code></p> <p><code>import</code> <code>torchvision.datasets as datasets</code></p> <p><code>class</code> <code>Net(nn.Module):</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;</code><code>def</code> <code>__init__(</code><code>self</code><code>):</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>super</code><code>(Net, </code><code>self</code><code>).__init__()</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>self</code><code>.conv1 </code><code>=</code> <code>nn.Conv2d(</code><code>3</code><code>, </code><code>6</code><code>, </code><code>5</code><code>)</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>self</code><code>.pool </code><code>=</code> <code>nn.MaxPool2d(</code><code>2</code><code>, </code><code>2</code><code>)</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>self</code><code>.conv2 </code><code>=</code> <code>nn.Conv2d(</code><code>6</code><code>, </code><code>16</code><code>, </code><code>5</code><code>)</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>self</code><code>.fc1 </code><code>=</code> <code>nn.Linear(</code><code>16</code> <code>*</code> <code>5</code> <code>*</code> <code>5</code><code>, </code><code>120</code><code>)</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>self</code><code>.fc2 </code><code>=</code> <code>nn.Linear(</code><code>120</code><code>, </code><code>84</code><code>)</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>self</code><code>.fc3 </code><code>=</code> <code>nn.Linear(</code><code>84</code><code>, </code><code>10</code><code>)</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;</code><code>def</code> <code>forward(</code><code>self</code><code>, x):</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>x </code><code>=</code> <code>self</code><code>.pool(F.relu(</code><code>self</code><code>.conv1(x)))</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>x </code><code>=</code> <code>self</code><code>.pool(F.relu(</code><code>self</code><code>.conv2(x)))</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>x </code><code>=</code> <code>x.view(</code><code>-</code><code>1</code><code>, </code><code>16</code> <code>*</code> <code>5</code> <code>*</code> <code>5</code><code>)</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>x </code><code>=</code> <code>F.relu(</code><code>self</code><code>.fc1(x))</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>x </code><code>=</code> <code>F.relu(</code><code>self</code><code>.fc2(x))</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>x </code><code>=</code> <code>self</code><code>.fc3(x)</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>return</code> <code>x</code></p> <p><code>net </code><code>=</code> <code>Net()</code></p> <p><code>PATH </code><code>=</code> <code>'./model.pth'</code></p> <p><code>torch.save(net.state_dict(), PATH)</code></p> <p><code>model </code><code>=</code> <code>Net()</code></p> <p><code>model.load_state_dict(torch.load(PATH))</code></p> <p><code>model.</code><code>eval</code><code>()</code></p> </td> 
  </tr> 
 </tbody> 
</table>

### 3. 调整输入和输出节点 ###

现在需要定义输入和输出节点，这些节点由导出的模型中的张量名称表示。将使用PyTorch内置的函数`torch.onnx.export()`来将模型转换为ONNX格式。下面的代码片段说明如何找到输入和输出节点，然后传递给该函数：

<table> 
 <tbody> 
  <tr> 
   <td> <p>1</p> <p>2</p> <p>3</p> <p>4</p> <p>5</p> <p>6</p> </td> 
   <td> <p><code>input_names </code><code>=</code> <code>[</code><code>"input"</code><code>]</code></p> <p><code>output_names </code><code>=</code> <code>[</code><code>"output"</code><code>]</code></p> <p><code>dummy_input </code><code>=</code> <code>torch.randn(batch_size, input_channel_size, input_height, input_width)</code></p> <p><code># Export the model</code></p> <p><code>torch.onnx.export(model, dummy_input, </code><code>"model.onnx"</code><code>, verbose</code><code>=</code><code>True</code><code>, </code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;</code><code>input_names</code><code>=</code><code>input_names, output_names</code><code>=</code><code>output_names)</code></p> </td> 
  </tr> 
 </tbody> 
</table>

### 4. 运行转换程序 ###

运行上述程序时可能遇到错误信息，其中包括一些与节点的名称和形状相关的警告，甚至还有Python版本、库、路径等信息。在处理完这些错误后，就可以转换PyTorch模型并立即获得ONNX模型了。输出ONNX模型的文件名是`model.onnx`。

### 5. 使用后端框架测试ONNX模型 ###

现在，使用ONNX模型检查一下是否成功地将其从PyTorch导出到ONNX，可以使用TensorFlow或Caffe2进行验证。以下是一个简单的示例，演示如何使用TensorFlow来加载和运行该模型：

<table> 
 <tbody> 
  <tr> 
   <td> <p>1</p> <p>2</p> <p>3</p> <p>4</p> <p>5</p> <p>6</p> <p>7</p> <p>8</p> </td> 
   <td> <p><code>import</code> <code>onnxruntime as rt</code></p> <p><code>import</code> <code>numpy as np</code></p> <p><code>sess </code><code>=</code> <code>rt.InferenceSession(</code><code>'model.onnx'</code><code>)</code></p> <p><code>input_name </code><code>=</code> <code>sess.get_inputs()[</code><code>0</code><code>].name</code></p> <p><code>output_name </code><code>=</code> <code>sess.get_outputs()[</code><code>0</code><code>].name</code></p> <p><code>np.random.seed(</code><code>123</code><code>)</code></p> <p><code>X </code><code>=</code> <code>np.random.randn(batch_size, input_channel_size, input_height, input_width).astype(np.float32)</code></p> <p><code>res </code><code>=</code> <code>sess.run([output_name], {input_name: X})</code></p> </td> 
  </tr> 
 </tbody> 
</table>

这应该可以顺利地运行，并且输出与原始PyTorch模型具有相同的形状（和数值）。

### 6. 核对结果 ###

最好的方法是比较PyTorch模型与ONNX模型在不同框架中推理的结果。如果结果完全匹配，则几乎可以肯定地说PyTorch到ONNX转换已经成功。以下是通过PyTorch和ONNX检查模型推理结果的一个小程序：

<table> 
 <tbody> 
  <tr> 
   <td> <p>1</p> <p>2</p> <p>3</p> <p>4</p> <p>5</p> <p>6</p> <p>7</p> <p>8</p> <p>9</p> </td> 
   <td> <p><code># Test the model with PyTorch</code></p> <p><code>model.</code><code>eval</code><code>()</code></p> <p><code>with torch.no_grad():</code></p> <p><code>&nbsp;&nbsp;&nbsp;&nbsp;</code><code>Y </code><code>=</code> <code>model(torch.from_numpy(X)).numpy()</code></p> <p><code># Test the ONNX model with ONNX Runtime</code></p> <p><code>sess </code><code>=</code> <code>rt.InferenceSession(</code><code>'model.onnx'</code><code>)</code></p> <p><code>res </code><code>=</code> <code>sess.run(</code><code>None</code><code>, {input_name: X})[</code><code>0</code><code>]</code></p> <p><code># Compare the results</code></p> <p><code>np.testing.assert_allclose(Y, res, rtol</code><code>=</code><code>1e</code><code>-</code><code>6</code><code>, atol</code><code>=</code><code>1e</code><code>-</code><code>6</code><code>)</code></p> </td> 
  </tr> 
 </tbody> 
</table>

以上就是PyTorch模型转换为ONNX格式的详细内容，希望可以帮到你！

转自：[微点阅读][Link 1] [https://www.weidianyuedu.com][Link 1]

[d1f5eadc79bdaef445efcee8564f8fbb.jpeg]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/13/ec6cfe192a404a45983bf2276a59c5be.png
[Link 1]: https://www.weidianyuedu.com

发表评论取消回复

表情：

评论列表（有 0 条评论，79人围观）

还没有评论，来说两句吧...

相关阅读

相关 PyTorch模型转换为ONNX格式实现过程浅谈

![d1f5eadc79bdaef445efcee8564f8fbb.jpeg][] 本文主要为大家介绍了PyTorch模型转换为ONNX格式实现过程详解，有需要的朋友可以借

浅浅的花香味﹌/ 2024年03月24日 18:47/ 0 赞/ 80 阅读

相关【yolov5】pytorch模型导出为onnx模型

> 博主想拿官网的yolov5训练好pt模型，然后转换成rknn模型，然后在瑞芯微开发板上调用模型检测。但是官网的版本对npu不友好，所以采用改进结构的版本： > > 1.

港控/mmm°/ 2023年10月02日 21:51/ 0 赞/ 18 阅读

相关转换onnx_PyTorch 到 ONNX 到 Tensorflow 到 TFLite的坑

.view() 在 Pytorch 中，经常会对某个具体的 Tensor 做变换的处理。但在 PyTorch 转换到 ONNX 的过程中，例如 fea = fe

电玩女神/ 2023年01月06日 01:20/ 0 赞/ 510 阅读

相关 pytorch转换onnx，再转换caffe 测试caffe, pytorch 模型结果是否一致

def pytorch_out(input): model = model_res() model.eval input = inpu

小灰灰/ 2023年01月03日 04:47/ 0 赞/ 191 阅读

相关 pytorch转换onnx,测试onnx,pytorch 模型结果是否一致

def to_onnx(): dummy_input = torch.randn(1, 3, 112, 112, dtype=torch.float)

小灰灰/ 2023年01月03日 04:16/ 0 赞/ 451 阅读

相关 Pytorch预训练模型文件转化为onnx格式

代码如下 coding: utf-8 import torch from model.wacnet import WACNet 导入自己构建好的模型

偏执的太偏执、/ 2022年12月11日 12:24/ 0 赞/ 319 阅读

相关 pytorch模型转ONNX模型

转载请注明作者和出处： [http://blog.csdn.net/john\_bh/][http_blog.csdn.net_john_bh] 文章目录

刺骨的言语ヽ痛彻心扉/ 2022年11月29日 05:56/ 0 赞/ 245 阅读

相关 Pytorch DataParallel多卡训练模型导出onnx模型

Pytorch模型转换到onnx模型代码如下： import torch import torch.nn as nn import torch.onn

我不是女神ヾ/ 2022年11月14日 13:27/ 0 赞/ 307 阅读

相关 Pytorch模型(.pth)转onnx模型(.onnx)

简介 Open Neural Network Exchange（ONNX，开放神经网络交换）格式，是一个用于表示深度学习模型的标准，可使模型在不同框架之间进行转移。 ON

Myth丶恋晨/ 2022年09月11日 02:25/ 0 赞/ 370 阅读

相关 Pytorch模型转onnx、caffe、TorchScript通用代码，onnx模型推理

[https://blog.csdn.net/zengwubbb/article/details/108696616?spm=1001.2014.3001.5501][http

待我称王封你为后i/ 2022年09月10日 12:22/ 0 赞/ 268 阅读