PyTorch 模型 onnx 文件导出及调用详情

介绍：

PyTorch是一个基于Python的科学计算库，它有诸多优异的特性，其中一个重要的特性是它的高效特定GPU加速的张量计算（tensor computation）操作。PyTorch 1.0版本（2018年12月）已经发布，包括了对多平台、多端到端场景的支持，同时完善了跨平台支持。

我们可以使用PyTorch训练模型，然后将训练好的模型导出为ONNX（Open Neural Network Exchange）格式，使用ONNX Runtime获得较快的预测速度。此外，我们还可以利用ONNX格式将训练好的模型轻松地部署到云端、移动设备等多个平台中。

在本文中，我们将介绍如何使用PyTorch导出ONNX模型，以及如何使用ONNX Runtime加载和运行模型。此外，还会提供两个简单的PyTorch模型示例以及相关代码，以帮助读者更好地理解PyTorch到ONNX的转换和操作。

Part 1：PyTorch 模型 onnx 文件导出

首先，安装ONNX包。使用以下命令进行安装：

pip install onnx

然后，让我们创建一个简单的模型。在本文的示例中，我们将创建一个简单的线性模型，该模型可以接受1个特征，并将其乘以一个数字，然后将两个数字相加，即：

import torch.nn as nn

class MyModel(nn.Module):
  def __init__(self):
    super(MyModel, self).__init__()
    self.fcn = nn.Linear(1, 1)
    self.fcn.weight.data.fill_(2.)
    self.fcn.bias.data.fill_(1.)

  def forward(self, x):
    x = self.fcn(x)
    return x

接下来，调用以下代码行将模型计算图从PyTorch转换为ONNX格式：

import torch.onnx
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--output', type=str, help='output model name', default='model.onnx')
args = parser.parse_args()

if __name__ == '__main__':
  model = MyModel()
  model.train(False)
  input_names = ["input"]
  output_names = ["output"]
  x = torch.randn(1, 1, 1, 1, requires_grad=True)
  torch.onnx.export(model,
                    x,
                    args.output,
                    input_names=input_names,
                    output_names=output_names)

该代码在训练模式下构建了PyTorch模型，并将其转换为evaluation模式。接下来，创建一个1x1x1x1的张量x，它将作为模型输入，用torch.onnx.export将模型转换为ONNX格式。输出的onnx文件将被存储在“filename”中。在我们的例子中，我们将文件名称设置为“model.onnx”。

Part 2：ONNX模型加载、运行及调用

现在，我们已经创建并导出了我们的模型，我们可以使用ONNX Runtime来加载和运行模型，使用以下两行代码来实现这个功能：

import onnxruntime as ort

model = ort.InferenceSession("model.onnx")

第一行导入ONNX Runtime库，第二行调用ort.InferenceSession方法，加载我们的onnx模型。在我们的示例中，我们将模型名称设置为“model.onnx”。

现在，我们可以使用以下代码来运行我们的模型并对其进行评估：

import numpy as np

x = np.array([2.], dtype=np.float32)
pred = model.run(["output"], {"input": x})
print(pred)

我们在此处指定模型的输入数据x并使用model.run进行评估。"output"和"input"是分别指定模型输出和输入的名字，并且我们将x作为一个字典传递给方法。该方法返回一个包含模型输出的字典。在我们的示例中，我们的模型输出第一层返回的值为[5.]，因为输入值为2，第一层返回的是4，加上偏置表示为：4 + 1 = 5。

至此，我们的ONNX模型创建并执行完毕。

Part 3：简单示例

以下是PyTorch中的两个简单示例，可以使读者更好地了解如何使用PyTorch导出ONNX模型并使用ONNX Runtime加载和运行模型。

示例1：线性模型

import torch.nn as nn
import torch.onnx
import argparse

class MyModel(nn.Module):
  def __init__(self):
    super(MyModel, self).__init__()
    self.fcn = nn.Linear(1, 1)
    self.fcn.weight.data.fill_(2.)
    self.fcn.bias.data.fill_(1.)

  def forward(self, x):
    x = self.fcn(x)
    return x

if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument('--output', type=str, help='output model name', default='model.onnx')
  args = parser.parse_args()

  model = MyModel()
  model.train(False)
  input_names = ["input"]
  output_names = ["output"]
  x = torch.randn(1, 1, 1, 1, requires_grad=True)
  torch.onnx.export(model,
                    x,
                    args.output,
                    input_names=input_names,
                    output_names=output_names)

该示例创建一个简单线性模型，我们将该模型导出为onnx格式。

示例2：卷积神经网络

import torch.nn as nn
import torch.onnx
import argparse

class MyModel(nn.Module):
  def __init__(self):
    super(MyModel, self).__init__()
    self.layer1 = nn.Sequential(
      nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=False),
      nn.BatchNorm2d(64),
      nn.ReLU(inplace=True),
      nn.MaxPool2d(kernel_size=2, stride=2)
    )
    self.fc = nn.Linear(64 * 14 * 14, 10)

  def forward(self, x):
    x = self.layer1(x)
    x = x.view(x.size(0), -1)
    x = self.fc(x)
    return x

if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument('--output', type=str, help='output model name', default='model.onnx')
  args = parser.parse_args()

  model = MyModel()
  model.train(False)
  input_names = ["input"]
  output_names = ["output"]
  x = torch.randn(1, 1, 28, 28, requires_grad=True)
  torch.onnx.export(model,
                    x,
                    args.output,
                    input_names=input_names,
                    output_names=output_names)

该示例创建了一个包含卷积神经网络的模型。我们将该模型导出为onnx格式。

这两个示例代码中的onnx格式文件导出和技术细节都已经在第一部分（Part 1）中进行了解释。在这里，我们主要关注如何导出、加载和运行onnx格式的PyTorch模型，以及如何将其用于新的应用程序、算法和技术中。

总结

本文介绍了如何使用PyTorch导出ONNX模型，并使用ONNX Runtime加载和运行模型，以及提供了两个简单的PyTorch模型示例。我们希望这篇文章能够帮助读者更好地转换和操作PyTorch到ONNX格式的模型。