[ONNX] Convert TorchScript model to Onnx model

pytorch 모델을 onnx모델로 변환하는 방법

 

pytorch 모델을 onnx모델로 변환하는 방법은 아래의 링크를 보면 설명이 되어 있다.

https://shuka.tistory.com/26

[Pytorch] pytorch모델을 onnx모델로 변환 방법

 

 

TorchScript모델을 onnx모델로 변환 중 발생한 Errror

TorchScript모델을 onnx모델로 변환하기 위해서는 위의 링크에서 얘기한 내용과 같이 export를 쓰면 되는데 내 경우에는 밑에와 같이 두 가지의 error가 발생했다.

 

1. TypeError: _convolution() missing 1 required positional arguement: 'allow_tf32'

 

다음과 같은 error를 찾아보다가 마땅한 수정 방법을 못찾아서 그냥 코드를 수정했다.

(내가 수정한 방법이 항상 정답은 아니고, 해결이 안되는 경우도 나올 수 있다.)

 

위의 오류는 torch.onnx.symbolic_opset9.py에 있는 _convoludion() 함수에서 밑의 그림과 같이 allow_tf32를 False처리해 주면 해결이 됐다.

 

 

2. AssertionError: example_outputs must be provided when exporting a ScriptModule or ScriptFunction.

해당 error는 export해줄 때 example_outputs를 설정해 달라는 메뉴다.

save_onnx_name = 'test.onnx'
dummy = torch.randn(1, 3, 256, 256).cuda()
torch.onnx.export(model, dummy, save_onnx_name, example_outputs=model(dummy))

 

 

Convert TorchScript model to onnx model

아래의 code는 TorchScript model을 load 하고 onnx로 export 하는 전체 code이다.

 

import torch
import onnx


model = torch.jit.load('hand.pts', map_location=torch.device('cuda'))
model.eval()
dummy = torch.randn(1, 3, 256, 256).cuda()

input_name = ['input_name']
output_name = ['output_name']
onnx_model = onnx.load("test.onnx")
torch.onnx.export(model, dummy, 'test.onnx', example_outputs=model(dummy), input_names=input_name, output_names=output_name)

 

Test Converted Onnx model

아래의 code는 변환된 onnx model에 대해 test를 해볼 수 있는 code이다.

import onnxruntime as ort

ort_session = ort.InferenceSession('test.onnx')
outputs = ort_session.run(None, {'input_name': to_numpy(dummy)})
print('outputs: ', outputs)