欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

pytorch模型转onnx

程序员文章站 2022-07-13 11:38:30
...

参考: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")