广州网站建设便宜,山西网络广告推广,网站怎么做可以增加点击率,免费申请版权前言
本文主要介绍如何将PyTorch模型转换为ONNX模型#xff0c;为后面的模型部署做准备。转换后的xxx.onnx模型#xff0c;进行加载和测试。最后介绍使用Netron#xff0c;可视化ONNX模型#xff0c;看一下网络结构#xff1b;查看使用了那些算子#xff0c;以便开发部署…前言
本文主要介绍如何将PyTorch模型转换为ONNX模型为后面的模型部署做准备。转换后的xxx.onnx模型进行加载和测试。最后介绍使用Netron可视化ONNX模型看一下网络结构查看使用了那些算子以便开发部署。 目录
前言
一、PyTorch模型转ONNX模型
1.1 转换为ONNX模型且加载权重
1.2 转换为ONNX模型但不加载权重
1.3 torch.onnx.export() 函数
二、加载ONNX模型
三、可视化ONNX模型 一、PyTorch模型转ONNX模型
将PyTorch模型转换为ONNX模型通常是使用torch.onnx.export( )函数来转换的基本的思路是
加载PyTorch模型可以选择只加载模型结构也可以选择加载模型结构和权重。然后定义PyTorch模型的输入维度比如(1, 3, 224, 224)这是一个三通道的彩色图分辨率为224x224。最后使用torch.onnx.export( )函数来转换生产xxx.onnx模型。
下面有一个简单的例子
import torch
import torch.onnx# 加载 PyTorch 模型
model ...# 设置模型输入包括通道数分辨率等
dummy_input torch.randn(1, 3, 224, 224, devicecpu)# 转换为ONNX模型
torch.onnx.export(model, dummy_input, model.onnx, export_paramsTrue)1.1 转换为ONNX模型且加载权重
这里举一个resnet18的例子基本思路是
首先加载了一个预训练的 ResNet18 模型然后将其设置为评估模式。接下来定义一个与模型输入张量形状相同的输入张量并使用 torch.randn() 函数生成了一个随机张量。最后使用 onnx.export() 函数将 PyTorch 模型转换为 ONNX 格式并将其保存到指定的输出文件中。
程序如下
import torch
import torchvision.models as models# 加载预训练的 ResNet18 模型
model models.resnet18(pretrainedTrue)# 将模型设置为评估模式
model.eval()# 定义输入张量需要与模型的输入张量形状相同
input_shape (1, 3, 224, 224)
x torch.randn(input_shape)# 需要指定输入张量输出文件路径和运行设备
# 默认情况下输出张量的名称将基于模型中的名称自动分配
device torch.device(cuda if torch.cuda.is_available() else cpu)# 将 PyTorch 模型转换为 ONNX 格式
output_file resnet18.onnx
torch.onnx.export(model, x.to(device), output_file, export_paramsTrue)1.2 转换为ONNX模型但不加载权重
举一个resnet18的例子基本思路是
首先加载了一个预训练的 ResNet18 模型然后使用 onnx.export() 函数将 PyTorch 模型转换为 ONNX 格式指定参数do_constant_foldingFalse不加载模型的权重。
import torch
import torchvision.models as models# 加载 PyTorch 模型
model models.resnet18()# 将模型转换为 ONNX 格式但不加载权重
dummy_input torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, resnet18.onnx, do_constant_foldingFalse)下面构建一个简单网络结构并转换为ONNX
import torch
import torchvision
import numpy as np# 定义一个简单的PyTorch 模型
class MyModel(torch.nn.Module):def __init__(self):super(MyModel, self).__init__()self.conv1 torch.nn.Conv2d(3, 32, kernel_size3, stride1, padding1)self.relu torch.nn.ReLU()self.maxpool torch.nn.MaxPool2d(kernel_size2, stride2)self.conv2 torch.nn.Conv2d(32, 64, kernel_size3, stride1, padding1)self.flatten torch.nn.Flatten()self.fc1 torch.nn.Linear(64 * 8 * 8, 10)def forward(self, x):x self.conv1(x)x self.relu(x)x self.maxpool(x)x self.conv2(x)x self.relu(x)x self.maxpool(x)x self.flatten(x)x self.fc1(x)return x# 创建模型实例
model MyModel()# 指定模型输入尺寸
dummy_input torch.randn(1, 3, 32, 32)# 将PyTorch模型转为ONNX模型
torch.onnx.export(model, dummy_input, mymodel.onnx, do_constant_foldingFalse) 1.3 torch.onnx.export() 函数
看一下这个函数的参数
torch.onnx.export(model, args, f, export_paramsTrue, opset_version10, do_constant_foldingTrue, input_names[input], output_names[output], dynamic_axesNone, verboseFalse, example_outputsNone, keep_initializers_as_inputsNone) model需要导出的 PyTorch 模型argsPyTorch模型输入数据的尺寸指定通道数、长和宽。可以是单个 Tensor 或元组也可以是元组列表。f导出的 ONNX 文件路径和名称mymodel.onnx。export_params是否导出模型参数。如果设置为 False则不导出模型参数。opset_version导出的 ONNX 版本。默认值为 10。do_constant_folding是否对模型进行常量折叠。如果设置为 True不加载模型的权重。input_names模型输入数据的名称。默认为 input。output_names模型输出数据的名称。默认为 output。dynamic_axes动态轴的列表允许在导出的 ONNX 模型中创建变化的维度。verbose是否输出详细的导出信息。example_outputs用于确定导出 ONNX 模型输出形状的样本输出。keep_initializers_as_inputs是否将模型的初始化器作为输入导出。如果设置为 True则模型初始化器将被作为输入的一部分导出。下面是只是一个常用的模板
import torch.onnx # 转为ONNX
def Convert_ONNX(model): # 设置模型为推理模式model.eval() # 设置模型输入的尺寸dummy_input torch.randn(1, input_size, requires_gradTrue) # 导出ONNX模型 torch.onnx.export(model, # model being run dummy_input, # model input (or a tuple for multiple inputs) xxx.onnx, # where to save the model export_paramsTrue, # store the trained parameter weights inside the model file opset_version10, # the ONNX version to export the model to do_constant_foldingTrue, # whether to execute constant folding for optimization input_names [modelInput], # the models input names output_names [modelOutput], # the models output names dynamic_axes{modelInput : {0 : batch_size}, # variable length axes modelOutput : {0 : batch_size}}) print( ) print(Model has been converted to ONNX)if __name__ __main__: # 构建模型并训练# xxxxxxxxxxxx# 测试模型精度#testAccuracy() # 加载模型结构与权重model Network() path myFirstModel.pth model.load_state_dict(torch.load(path)) # 转换为ONNX Convert_ONNX(model) 二、加载ONNX模型
加载ONNX模型通常需要用到ONNX、ONNX Runtime所以需要先安装。
pip install onnx
pip install onnxruntime
加载ONNX模型可以使用ONNX Runtime库以下是一个加载ONNX模型的示例代码
import onnxruntime as ort# 加载 ONNX 模型
ort_session ort.InferenceSession(model.onnx)# 准备输入信息
input_info ort_session.get_inputs()[0]
input_name input_info.name
input_shape input_info.shape
input_type input_info.type# 运行ONNX模型
outputs ort_session.run(input_name, input_data)# 获取输出信息
output_info ort_session.get_outputs()[0]
output_name output_info.name
output_shape output_info.shape
output_data outputs[0]print(outputs:, outputs)
print(output_info :, output_info )
print(output_name :, output_name )
print(output_shape :, output_shape )
print(output_data :, output_data ) 以下是一个示例程序将 resnet18 模型从 PyTorch 转换为 ONNX 格式然后加载和测试 ONNX 模型的过程:
import torch
import torchvision.models as models
import onnx
import onnxruntime# 加载 PyTorch 模型
model models.resnet18(pretrainedTrue)
model.eval()# 定义输入和输出张量的名称和形状
input_names [input]
output_names [output]
batch_size 1
input_shape (batch_size, 3, 224, 224)
output_shape (batch_size, 1000)# 将 PyTorch 模型转换为 ONNX 格式
torch.onnx.export(model, # 要转换的 PyTorch 模型torch.randn(input_shape), # 模型输入的随机张量resnet18.onnx, # 保存的 ONNX 模型的文件名input_namesinput_names, # 输入张量的名称output_namesoutput_names, # 输出张量的名称dynamic_axes{input_names[0]: {0: batch_size}, output_names[0]: {0: batch_size}} # 动态轴即输入和输出张量可以具有不同的批次大小
)# 加载 ONNX 模型
onnx_model onnx.load(resnet18.onnx)
onnx_model_graph onnx_model.graph
onnx_session onnxruntime.InferenceSession(onnx_model.SerializeToString())# 使用随机张量测试 ONNX 模型
x torch.randn(input_shape).numpy()
onnx_output onnx_session.run(output_names, {input_names[0]: x})[0]print(fPyTorch output: {model(torch.from_numpy(x)).detach().numpy()[0, :5]})
print(fONNX output: {onnx_output[0, :5]})上述代码中首先加载预训练的 resnet18 模型并定义了输入和输出张量的名称和形状。
然后使用 torch.onnx.export() 函数将模型转换为 ONNX 格式并保存为 resnet18.onnx 文件。
接着使用 onnxruntime.InferenceSession() 函数加载 ONNX 模型并使用随机张量进行测试。
最后将 PyTorch 模型和 ONNX 模型的输出进行比较以确保它们具有相似的输出。 三、可视化ONNX模型
使用Netron可视化ONNX模型看一下网络结构查看使用了那些算子以便开发部署。
这里简单介绍一下 Netron是一个轻量级、跨平台的模型可视化工具支持多种深度学习框架的模型可视化包括TensorFlow、PyTorch、ONNX、Keras、Caffe等等。它提供了可视化网络结构、层次关系、输出尺寸、权重等信息并且可以通过鼠标移动和缩放来浏览模型。Netron还支持模型的导出和导入方便模型的分享和交流。 Netron的网页在线版本直接在网页中打开和查看ONNX模型Netron
开源地址GitHub - lutzroeder/netron: Visualizer for neural network, deep learning, and machine learning models 支持多种操作系统
macOS: Download
Linux: Download
Windows: Download
Browser: Start
Python Server: Run pip install netron and netron [FILE] or netron.start([FILE]). 下面是可视化模型截图 还能查看某个节点运算操作的信息比如下面MaxPool点击一下能看到使用的3x3的池化核是否有填充pads步长strides等参数。 分享完毕~
