onnx文件转pytorch pt模型文件

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