torch与onnx转换
torch与onnx转换
一、pytorch转onnx
pytorch官方已经提供了支持。
import torch
dummy_input = torch.randn(1, 3, 480, 640)
torch.onnx.export(model, dummy_input, 'net_640x480.onnx', export_params=True, verbose=False, training=False, input_names=['Input_1'], output_names=['Concat_211', 'Reshape_235'], opset_version=11, dynamic_axes={"Input_1":{0:"batch_size"}, "Concat_211":{0:"batch_size"}, "Reshape_235":{0:"batch_size"}})
导出onnx后,可以使用onnx-simplify去掉模型的一些冗余操作:
from onnxsim import simplify
import onnx
onnx_model = onnx.load("net_640x480.onnx")
model_simp, check = simplify(onnx_model, input_shapes={'Input_1':[1, 3, 480, 640]})
onnx.save(model_simp, 'net_640x480_simplify.onnx')
二、onnx转pytorch
可以使用第三方工具onnx2torch:
https://github.com/fumihwh/onnx-pytorch.git
查看代码:
也可以参考:
https://github.com/ToriML/onnx2pytorch/tree/master/onnx2pytorch
三、其它
可以使用在线转换工具进行模型的转换:
![转换工具 转换工具](https://gitee.com/bairuiworld/storywriter_image/raw/master/story/1638412540667.png)
转换工具