pytorch模型转onnx

CV-deeplearning

已于 2023-02-20 13:07:32 修改

阅读量547

点赞数

分类专栏：模型部署文章标签：人工智能

于 2020-09-27 09:51:19 首次发布

本文链接：https://blog.csdn.net/Guo_Python/article/details/108821599

版权

模型部署专栏收录该内容

15 篇文章 1 订阅

订阅专栏

参考：pytorch官网

本文根据官网的例子，把一个mobilenetV2的分类模型转化为onnx

主要分为一下几步：

1. 导入模型，加载权重参数，将模型设置为eval模式。

2. 构造随机的输入，输入的channel, height, weight要和训练时候一致。

3. 导出onnx。

import io
import torch
import torch.onnx
from MobileNetV2 import mobilenet_v2


torch_model = mobilenet_v2()

state_dict = torch.load("checkpoint/mobilenet-v2_0.pth", map_location='cuda:0')

torch_model.load_state_dict(state_dict)
torch_model.eval()

batch_size = 1
x = torch.randn(batch_size, 3, 128, 128, requires_grad=True)

torch.onnx.export(torch_model, x, "mobilenet_v2.onnx")

onnx可视化编辑请参考：

onnx-modifier：ONNX可视化编辑_onnx modifier_歌你一下的博客-CSDN博客