导入模型和运行模型预测

一、导入已有模型

这一小节主要介绍如何加载一个模型以及它的持久参数状态和推理模型预测。

%matplotlib inlineimport torchimport onnxruntimefrom torch import nnimport numpy as npimport torch.onnx as onnximport torchvision.models as modelsfrom torchvision import datasetsfrom torchvision.transforms import ToTensor

为了载入模型，我们需要预先定义一个模型类，里面包含用于训练模型的神经网络的状态和参数。

# 验证是否安装成功onnxruntime.get_device()

'GPU'

class NeuralNetwork(nn.Module):    def __init__(self):        super(NeuralNetwork, self).__init__()        self.flatten = nn.Flatten()        self.linear_relu_stack = nn.Sequential(            # 下面的各层顺序和参数个数都要和将导入的保持一致            nn.Linear(28*28, 512),            nn.ReLU(),            nn.Linear(512, 512),            nn.ReLU(),            nn.Linear(512, 10),            nn.ReLU()        )    def forward(self, x):        x = self.flatten(x)        logits = self.linear_relu_stack(x)        return logits

在载入模型权重时，我们首先要实例化定义的模型类，因为这个类定一个nn的结构。下面，我们使用load_state_dict()方法来载入模型的参数

# 先实例化预定义的网络模型，来为将要导入的模型定好"骨架"model = NeuralNetwork()# 将模型的参数导入，“丰满血肉”model.load_state_dict(torch.load('data/model.pth'))model.eval()

NeuralNetwork(  (flatten): Flatten(start_dim=1, end_dim=-1)  (linear_relu_stack): Sequential(    (0): Linear(in_features=784, out_features=512, bias=True)    (1): ReLU()    (2): Linear(in_features=512, out_features=512, bias=True)    (3): ReLU()    (4): Linear(in_features=512, out_features=10, bias=True)    (5): ReLU()  ))

注意：

Dropout 是一种用于神经网络防止过拟合的正则化技术。 它通过在每次训练迭代中随机地设置神经元中的一小部分为 0 来阻止神经元共适应（co-adapting），Dropout 可以通过多种方式进行解读，比如从不同网络的指数数字中随机取样。

确保在进行推理之前调用 model.eval ()方法，将dropout和批量泛化层设置为评估模式。这是为了避免得到不同的推理结果

二、模型推断

一般来说，使一个运行在大量平台和编码语言上的模型最优化是很困难的。在所有框架和硬件的不同组合中最大化性能是非常耗时的。开放式神经网络交换(ONNX) 运行时为您提供了一次性训练的解决方案，并提高了对任何硬件、云或边缘设备的推理能力。

ONNX 是许多供应商支持的共享神经网络和其他机器学习模型的通用格式。您可以使用 ONNX 格式对其他编程语言和框架(如 Java、 JavaScript、 c # 和 ML.NET)上的模型进行推理

三、将模型导出到ONNX

PyTorch 还支持本地 ONNX 导出。然而，考虑到 PyTorch 执行图的动态特性，导出过程必须遍历执行图以生成持久的 ONNX 模型。出于这个原因，应该向导出流程传递一个具有适当大小的测试变量(在我们的例子中，我们将创建一个具有适当大小的虚拟零张量)。您可以从您的训练数据集的形状功能获得大小。例如，tensor.shape) :

input_image = torch.zeros((1,28,28))onnx_model = 'data/model.onnx'onnx.export(model, input_image, onnx_model)

我们将使用测试数据集作为样本数据，从 ONNX 模型进行推断以进行预测。

test_data = datasets.FashionMNIST(    root="data",    train=False,    download=True,    transform=ToTensor())classes = [    "T-shirt/top",    "Trouser",    "Pullover",    "Dress",    "Coat",    "Sandal",    "Shirt",    "Sneaker",    "Bag",    "Ankle boot",]# x, y 分别是测试集中的的第 1个测试样本的 Tensor 阵列和对应分类下标x, y = test_data[0][0], test_data[0][1]

x,y

(tensor([[[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,           0.0000, 0.0000, 0.0000, 0.0000],          ...          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,           0.0000, 0.0000, 0.0000, 0.0000]]]), 9)

我们需要使用 onnxruntime.InferenceSession 创建推理会话。要推断 onnx 模型，使用 run 并传递希望返回的输出列表(如果需要所有输出，请保持空白)和输入值的映射。结果是输出的列表。

# 验证是否安装成功onnxruntime.get_device()0

# 验证是否安装成功onnxruntime.get_device()1

# 验证是否安装成功onnxruntime.get_device()2

# 验证是否安装成功onnxruntime.get_device()3

可以看出，仍存在误差

原文：https://juejin.cn/post/7095208942459944974