一、torch.onnx.export()详细介绍

1.torch.onnx.export()

torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None, aten=False, export_raw_ir=False, operator_export_type=None, opset_version=None, _retain_param_name=True, do_constant_folding=False, example_outputs=None, strip_doc_string=True, dynamic_axes=None, keep_initializers_as_inputs=None)

2. 功能：

将pth模型转为onnx文件导出。

3.参数

model (torch.nn.Module) ：pth模型文件；
args (tuple of arguments) ：模型的输入, 模型的尺寸；
export_params (bool, default True) – 如果指定为True或默认, 参数也会被导出. 如果你要导出一个没训练过的就设为 False；
verbose (bool, default False) ：导出轨迹的调试描述；
training (bool, default False) ：在训练模式下导出模型。目前，ONNX导出的模型只是为了做推断，通常不需要将其设置为True；
input_names (list of strings, default empty list) ：onnx文件的输入名称；
output_names (list of strings, default empty list) ：onnx文件的输出名称；
opset_version：默认为9；
dynamic_axes – {‘input’ : {0 : ‘batch_size’}, ‘output’ : {0 : ‘batch_size’}}) 。

二、pth的保存方式

torch.save(model,'save_path')

torch.save(model,path) 会将model的参数、框架都保存到路径path中，但是在加载model的时候可能会因为包版本的不同报错，所以当保存所有模型参数时，需要将模型构造相关代码文件放在相同路径，否则在load的时候无法索引到model的框架。

torch.save(model.state_dict(),model_path)

建议：使用state_dict()模式保存model，torch.save(model.state_dict(),path)，这样保存为字典模式，可以直接load。

三、pth转onnx代码

1.使用torch.save(model,'save_path')方式保存

x = torch.randn(1, 3, 224, 224, device=device)

输入测试数据数据格式[batch, channl, height, width]

model.eval()

不启用 BatchNormalization 和 Dropout，保证BN和dropout不发生变化，pytorch框架会自动把BN和Dropout固定住，不会取平均，而是用训练好的值，不然的话，一旦test的batch_size过小，很容易就会被BN层影响结果。

注：一定要写上这句话，不然可能会影响onnx的输出结果，经验所知。

import torch
import torch.nn
import onnx# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = torch.device('cpu')model = torch.load('***.pth', map_location=device)
model.eval()input_names = ['input']
output_names = ['output']x = torch.randn(1, 3, 224, 224, device=device)torch.onnx.export(model, x, '***.onnx', input_names=input_names, output_names=output_names, verbose='True')

2.使用torch.save(model.state_dict(),model_path)方式保存

该方式保存需要提供网络结构文件。

import torch.onnx
import onnxruntime as ort
from model import Net# 创建.pth模型model = Net
# 加载权重
model_path = '***.pth'
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_statedict = torch.load(model_path, map_location=device)
model.load_state_dict(model_statedict)model.to(device)
model.eval()input_data = torch.randn(1, 3, 224, 224, device=device)# 转化为onnx模型
input_names = ['input']
output_names = ['output']torch.onnx.export(model, input_data, '***.onnx', opset_version=9, verbose=True, input_names=input_names, output_names = output_names)

pth文件转为onnx格式相关推荐

使用onnx包将pth文件转换为onnx文件
本文对比一下两种pth文件转为onnx的区别以及onnx文件在NETRON中的图只有参数的pth文件:cat_dog.pth 既有参数又有模型结构的pth文件:cat_dog_model_args. ...
pythoncsv格式_python实现csv格式文件转为asc格式文件的方法
一.背景描述 csv格式文件是一种类似于excel的文件格式 asc格式文件是一种可以用text打开的文本文件 csv转asc本来可以用arcgis顺利完成,但由于csv数据量太大(744万行),ar ...
python 读取csv文件转成字符串,python实现csv格式文件转为asc格式文件的方法
一.背景描述 csv格式文件是一种类似于excel的文件格式 asc格式文件是一种可以用text打开的文本文件 csv转asc本来可以用arcgis顺利完成,但由于csv数据量太大(744万行),ar ...
Python之ffmpeg：利用python编程基于ffmpeg将m4a格式音频文件转为mp3格式文件
Python之ffmpeg:利用python编程基于ffmpeg将m4a格式音频文件转为mp3格式文件目录利用python编程基于ffmpeg将m4a格式音频文件转为mp3格式文件 1.先下载ff ...
ppt生成eps文件_如何将AI/EPS格式文件转为ppt格式
平时见到很多AI或EPS文件都想直接拿来放到PPT上,对于很多图片来说,只需转PNG就可以导入到PPT了,但对于一些图表,在导入PPT的时候还想要编辑下,这个时候就要用到下面的将AI/EPS格式文件转 ...
批量将json文件转为jpg格式
批量将json文件转为jpg格式 1.简介 labelme软件自带有将json文件转为jpg的代码,不过只限单张. 位于./cli/json_to_dataset.py. 2.批量生成代码 impor ...
python3 将eps文件转为jpg格式
文章目录 1. 按 2. 准备 3. 代码 3.1. 保存成eps格式 3.2. 将eps文件转为jpg格式 1. 按用Turtle画的图无法直接保存成jpg格式的,只能先保存成eps,再将eps转 ...
python 使用字节流bytes格式读取文件转为int格式，再转为0,1字符串格式
python 使用字节流bytes格式读取文件转为int格式,再转为0,1字符串格式 with open('test.jpg', 'rb') as src:t = src.read(1) # 读进1B ...
批量将NC文件转为tif格式
批量将多年的NC文件转为tif格式进行处理(来源:https://www.geek-share.com/detail/2763962738.html) 所参考和借鉴的文章的链接如下: https:// ...
利用Python将WEBVTT格式的视频字幕文件转为SRT格式
1 WebVTT & SRT 格式 WebVTT字幕格式与SRT字幕格式主要区别在于时间格式的区分. 下面是一个WebVTT格式的字幕文件 WEBVTT1 00:00:20.000 --> ...

pth文件转为onnx格式