将pytorch训练保存的pth文件转为onnx文件,为后续模型部署做准备。
一、分类模型
import torch
import os
import timm
import argparse
from utils_net import Resnet
parser = argparse.ArgumentParser()
parser.add_argument("--pth_path", default='classify_model.pth')
parser.add_argument("--save_onnx_path", default='classify_model.onnx')
parser.add_argument("--input_width", default=416)
parser.add_argument("--input_height", default=416)
parser.add_argument("--input_channel", default=1)
parser.add_argument("--num_classes", default=6)
args = parser.parse_args()
def pth_to_onnx(pth_path, onnx_path, in_hig, in_wid, in_chal, num_cls):
if not onnx_path.endswith('.onnx'):
print('Warning! The onnx model name is not correct,\
please give a name that ends with \'.onnx\'!')
return 0
model = Resnet(num_classes=num_cls)
model.load_state_dict(torch.load(pth_path))
model.eval()
print(f'{pth_path} model loaded')
input_names = ['input']
output_names = ['output']
im = torch.rand(1, in_chal, in_hig, in_wid)
torch.onnx.export(model, im, onnx_path,
verbose=False,
input_names=input_names,
output_names=output_names)
print("Exporting .pth model to onnx model has been successful!")
print(f"Onnx model save as {onnx_path}")
if __name__ == '__main__':
pth_to_onnx(pth_path=args.pth_path,
onnx_path=args.save_onnx_path,
in_hig=args.input_height,
in_wid=args.input_width,
in_chal=args.input_channel,
num_cls=args.num_classes)
运行结果:
classify_model.pth model loaded
Exporting .pth model to onnx model has been successful!
Onnx model save as classify_model.onnx
Process finished with exit code 0
二、分割模型
import torch
import os
import argparse
from utils_net import seg_net
parser = argparse.ArgumentParser()
parser.add_argument("--pth_path", default='segment_model.pth')
parser.add_argument("--save_onnx_path", default='segment_model.onnx')
parser.add_argument("--input_width", default=416)
parser.add_argument("--input_height", default=416)
parser.add_argument("--input_channel", default=1)
parser.add_argument("--num_classes", default=4)
args = parser.parse_args()
def pth_to_onnx(pth_path, onnx_path, in_hig, in_wid, in_channel, num_cls):
if not onnx_path.endswith('.onnx'):
print('Warning! The onnx model name is not correct,\
please give a name that ends with \'.onnx\'!')
return 0
model = seg_net(in_channel=in_channel, num_cls=num_cls)
model.load_state_dict(torch.load(pth_path))
model.eval()
print(f'{pth_path} model loaded')
input_names = ['input']
output_names = ['output']
im = torch.rand(1, in_channel, in_hig, in_wid)
torch.onnx.export(model, im, onnx_path,
verbose=False,
input_names=input_names,
output_names=output_names,
opset_version=11)
print("Exporting .pth model to onnx model has been successful!")
print(f"Onnx model save as {onnx_path}")
if __name__ == '__main__':
pth_to_onnx(pth_path=args.pth_path,
onnx_path=args.save_onnx_path,
in_hig=args.input_height,
in_wid=args.input_width,
in_channel=args.input_channel,
num_cls=args.num_classes)
运行结果:
segment_model.pth model loaded
Exporting .pth model to onnx model has been successful!
Onnx model save as segment_model.onnx
Process finished with exit code 0
三、目标检测模型
在这里插入代码片
import torch
import onnx
import argparse
from utils_net import YoloBody
parser = argparse.ArgumentParser()
parser.add_argument("--pth_path", default='yolo.pth')
parser.add_argument("--save_onnx_path", default='yolo.onnx')
parser.add_argument("--input_width", default=416)
parser.add_argument("--input_height", default=416)
parser.add_argument("--num_classes", default=2)
parser.add_argument("--anchors_mask", default=[[6, 7, 8], [3, 4, 5], [0, 1, 2]])
args = parser.parse_args()
def pth_to_onnx(pth_path: str, save_onnx_path: str, num_cls: int,
in_hig: int, in_wid: int, anchor_mask: list,
opset_version: int = 12, simplify: bool = False):
"""
:param pth_path: pth文件文件
:param save_onnx_path: 准备保存的onnx路径
:param num_cls: 检测目标类别数
:param in_hig: 网络输入高度
:param in_wid: 网络输入宽度
:param anchor_mask: anchor宽高索引
:param opset_version: onnx算子集版本
:param simplify: 是否对模型进行简化
:return:保存onnx到指定路径
"""
# Build model, load weights
net = YoloBody(anchors_mask=anchor_mask,
num_classes=num_cls)
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# net.load_state_dict(torch.load(pth_path, map_location=device))
net.load_state_dict(torch.load(pth_path))
# print(next(net.parameters()).device)
net = net.eval()
print(f'{pth_path} model loaded')
im = torch.zeros(1, 3, in_hig, in_wid).to('cpu')
input_layer_names = ['images']
output_layer_names = ['output']
# Export the model
print(f'Starting export with onnx {onnx.__version__}.')
torch.onnx.export(net,
im,
f=save_onnx_path,
verbose=False,
opset_version=opset_version,
training=torch.onnx.TrainingMode.EVAL,
do_constant_folding=True,
input_names=input_layer_names,
output_names=output_layer_names,
dynamic_axes=None)
# Checks
model_onnx = onnx.load(save_onnx_path) # load onnx model
onnx.checker.check_model(model_onnx) # check onnx model
# Simplify onnx
if simplify:
import onnxsim
print(f'Simplifying with onnx-simplifier {onnxsim.__version__}.')
model_onnx, check = onnxsim.simplify(
model_onnx,
dynamic_input_shape=False,
input_shapes=None)
assert check, 'assert check failed'
onnx.save(model_onnx, save_onnx_path)
print('Onnx model save as {}'.format(save_onnx_path))
if __name__ == '__main__':
pth_to_onnx(pth_path=args.pth_path,
save_onnx_path=args.save_onnx_path,
num_cls=args.num_classes,
in_hig=args.input_height,
in_wid=args.input_width,
anchor_mask=args.anchors_mask)
运行结果:
yolo.pth model loaded
Starting export with onnx 1.11.0.
Onnx model save as yolo.onnx
Process finished with exit code 0