导出、校验pytorch为onnx模型

最新推荐文章于 2024-04-29 21:54:13 发布

Ceri

最新推荐文章于 2024-04-29 21:54:13 发布

阅读量299

点赞数

文章标签： pytorch python 深度学习

本文链接：https://blog.youkuaiyun.com/qq_38806886/article/details/130323067

版权

该代码示例展示了如何将一个预训练的PyTorchmobilenet_v3_large模型加载并调整，然后导出为ONNX格式。首先，模型的输出层被修改以适应特定的类别数。接着，使用torch.onnx.export函数将模型转换为ONNX模型，并进行简化。最后，使用ONNXRuntime测试转换后的模型，确保其能正确运行。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

以torchvision中mobilenet_v3_large模型为例

import torch
import onnx

import numpy as np
import torch.nn as nn
import onnxruntime as ort

from torchvision.models import *

torch_model = torch.load("***.pth") # pytorch模型加载
model = mobilenet_v3_large(pretrained=True)
num_ftrs = model.classifier[3].in_features
model.classifier[3] = nn.Linear(num_ftrs, 7) # 类别数量

model.load_state_dict(torch_model) 
dummy_input = torch.ones(1, 3, 112, 112)  # 输入图像大小

#测试pt模型
model.eval()
outputs = model(dummy_input)
print("pt_out:",outputs[0])

export_onnx_file = "0423***.onnx"          # 目的ONNX文件名
torch.onnx.export(model.eval(), dummy_input, export_onnx_file, verbose=False, input_names=["input"], output_names=["output"], opset_version=10)
print("======================== convert onnx Finished! .... ")

try:
    import onnxsim
    onnx_model = onnx.load(export_onnx_file)
    print('\nStarting to simplify ONNX...')
    sim_onnx_model, check = onnxsim.simplify(onnx_model)
    assert check, 'assert check failed'
except Exception as e:
    print(f'Simplifier failure: {e}')
onnx.save(sim_onnx_model,export_onnx_file)
print("Success to Simplify !!!")


#测试onnx
input_onnx = dummy_input.numpy().astype(np.float32)
ort_session = ort.InferenceSession('./0423***.onnx')
outs = (ort_session.run(None, {'input': input_onnx}))
print("onnx_out:",outs[0][0])