onnx文件转pytorch pt模型文件
从pytorch
格式转onnx
格式,官方有成熟的API;那么假如只有onnx
格式的模型文件,该怎样转回pytorch格式?
https://github.com/ENOT-AutoDL/onnx2torch提供了一种解决方法。
1.onnx2torch转换及测试
整体使用非常简单:
python
import torch
import onnxruntime
import numpy as np
onnx_file="model.onnx"
pt_file="model.pt"
DEVICE = 'cpu'
_input = torch.rand([1, 3, 512, 512]).to(DEVICE)
providers = ['CPUExecutionProvider']
session = onnxruntime.InferenceSession(onnx_file, providers=providers)
output_names = [x.name for x in session.get_outputs()]
onnx_output= session.run(output_names, {session.get_inputs()[0].name: _input.cpu().numpy()})
print(onnx_output[0].shape)
model=torch.load(pt_file)
model.to(DEVICE)
model.eval()
with torch.no_grad():
torch_output = model(_input)
print(torch_output.shape)
print(torch.allclose(torch.tensor(onnx_output[0]),torch_output, atol=1e-3))
打印结果可知,转换精度是很高的。
2.存在问题
转换成的torch
模型使用时batch
要与导出onnx
时使用的batch
一致,否则会出错。
有机会再研究如何修改。
Anyway,这已经提供了非常好的方案了。
参考文献
[1] https://github.com/ENOT-AutoDL/onnx2torch
[2] MMdnn