pytorch模型转onnx

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