pytorch基础6-加载并运行模型预测 & 总结
在本专题中,我们介绍了使用神经网络构建机器学习模型的关键概念,并使用PyTorch实现了这些概念。我们构建了一个图像识别模型,可以对图像进行分类,例如:T恤/上衣、裤子、套衫、连衣裙、外套、凉鞋、衬衫、运动鞋、包包和短靴。如何在CPU和GPU上使用张量如何管理、缩放和规范化数据集如何使用神经网络构建模型如何优化模型如何优化模型推理。
专题链接:https://blog.csdn.net/qq_33345365/category_12591348.html
本教程翻译自微软教程:https://learn.microsoft.com/en-us/training/paths/pytorch-fundamentals/
初次编辑:2024/3/4;最后编辑:2024/3/4
本教程第一篇:介绍pytorch基础和张量操作
本教程第二篇:介绍了数据集与归一化
本教程第三篇:介绍构建模型层的基本操作。
本教程第四篇:介绍自动微分相关知识。
本教程第五篇:介绍优化循环(optimization loop)相关知识。
本教程第六篇(即本博客内容):
- 介绍加载与运行模型预测的相关知识。
- 总结
另外本人还有pytorch CV相关的教程,见专题:
https://blog.csdn.net/qq_33345365/category_12578430.html
加载和运行模型预测
加载模型
本单元将介绍如何加载一个带有持久化参数状态和推断模型预测的模型。下面是本教程使用的类:
import torch
import onnxruntime
from torch import nn
import torch.onnx as onnx
import torchvision.models as models
from torchvision import datasets
from torchvision.transforms import ToTensor
为了加载模型,需要定义模型类,其中包含用于训练模型的神经网络的状态和参数。
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
在加载模型权重时,需要首先实例化模型类,因为类定义了网络的结构。接下来,使用load_state_dict()方法加载参数。
model = NeuralNetwork()
model.load_state_dict(torch.load('data/model.pth'))
model.eval()
**注意:**在进行推理之前,请确保调用
model.eval()方法,将dropout和批量归一化层(batch normalization layers)设置为评估模式(evaluation mode)。否则,您将看到不一致的推理结果。
模型推理
将神经网络模型优化的在各种平台和编程语言上运行是很困难的。在不同的框架和硬件组合中最大化性能非常耗时。**ONNX(Open Neural Network Exchange)**运行时为您提供了一种解决方案,您可以在任何硬件、云或边缘设备上训练一次并加速推理。
ONNX是一种通用格式,受多个供应商支持,用于共享神经网络和其他机器学习模型。您可以使用ONNX格式在其他编程语言和框架上进行推理,例如Java、JavaScript、C#和ML.NET。
输出模型到ONNX格式
PyTorch也具有原生的ONNX导出支持。然而,由于PyTorch执行图的动态性质,导出过程必须遍历执行图以产生持久化的ONNX模型。因此,在导出过程中应传入一个适当大小的测试变量(在我们的情况下,我们将创建一个正确大小的零张量作为虚拟数据。您可以使用shape函数在训练数据集上获取大小,如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 = test_data[0][0], test_data[0][1]
可以使用onnxruntime.InferenceSession创建一个推理会话。要推理ONNX模型,调用run并传入想要返回的输出列表(如果想要全部输出,则留空)和输入值的映射。结果是一个输出列表。
session = onnxruntime.InferenceSession(onnx_model, None)
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
result = session.run([output_name], {input_name: x.numpy()})
predicted, actual = classes[result[0][0].argmax(0)], classes[y]
print(f'Predicted: "{predicted}", Actual: "{actual}"')
输出是:
Predicted: "Ankle boot", Actual: "Ankle boot"
ONNX模型使您能够在不同的平台和不同的编程语言中运行推理。
代码汇总
import torch
import onnxruntime
from torch import nn
import torch.onnx as onnx
import torchvision.models as models
from torchvision import datasets
from torchvision.transforms import ToTensor
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
model = NeuralNetwork()
model.load_state_dict(torch.load('data/model1.pth'))
model.eval()
input_image = torch.zeros((1, 28, 28))
onnx_model = 'data/model.onnx'
onnx.export(model, input_image, onnx_model)
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 = test_data[0][0], test_data[0][1]
session = onnxruntime.InferenceSession(onnx_model, None)
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
result = session.run([output_name], {input_name: x.numpy()})
predicted, actual = classes[result[0][0].argmax(0)], classes[y]
print(f'Predicted: "{predicted}", Actual: "{actual}"')
总结
在本专题中,我们介绍了使用神经网络构建机器学习模型的关键概念,并使用PyTorch实现了这些概念。我们构建了一个图像识别模型,可以对图像进行分类,例如:T恤/上衣、裤子、套衫、连衣裙、外套、凉鞋、衬衫、运动鞋、包包和短靴。
您学习了以下关键领域:
- 如何在CPU和GPU上使用张量
- 如何管理、缩放和规范化数据集
- 如何使用神经网络构建模型
- 如何优化模型
- 如何优化模型推理
更多推荐

所有评论(0)