Torch model to ONNX model 변환 시 유의사항 (pytorch model과 ONNX 모델 결과가 다를 때)

pytorch model을 ONNX로 변환해서 사용하고자 할 때 onnx를 직접 쓸 수도 있지만 가장 편하게 쓰는 방법은 torch.onnx.export를 쓰는 방법이다. 

 

https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html

(optional) Exporting a Model from PyTorch to ONNX and Running it using ONNX Runtime — PyTorch Tutorials 2.4.0+cu121 documentat

 

# Input to the model
x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)
torch_out = torch_model(x)

# Export the model
torch.onnx.export(torch_model,               # model being run
                  x,                         # model input (or a tuple for multiple inputs)
                  "super_resolution.onnx",   # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=10,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input'],   # the model's input names
                  output_names = ['output'], # the model's output names
                  dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes
                                'output' : {0 : 'batch_size'}})

 

예제에서도 나와있듯이, 위와 같이 torch model과 dummy 입력을 같이 통과시키면 변환이 완료된다. 

 

근데 변환까지는 성공하는데 변환 결과 (inference 결과)가 pytorch model 전혀 다르게 나올 때가 있다.

(성능이 떨어지는 수준이 아니라 값이 잘못 나오거나 dimension이 잘못나오는 경우)

 

원인은 ONNX graph 최적화 과정에서 pytorch model의 일부를 삭제해버렸기 때문이다. 

 

https://netron.app/

Netron

위 페이지에 들어가 drag & drop 으로 ONNX 모델을 던져보면 설계한 pytorch model과 다른 부분이 있다.

random noise 입력 넣어줬을 때

inference에 사용할 입력 넣어줬을 때

 

위 예시는 같은 pytorch model을 변환했을 때 나온 두가지 ONNX graph인데, 완전히 구조가 바뀌어 있는 것을 볼 수 있다.

 

유의사항

원인은 다름 아니라 변환 시 사용한 입력에 있다. 입력은 실제 inference에 쓸 입력으로 넣어주는게 맞다.

x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)

 

예제에서는 아무런 random 값을 넣어도 되는 것처럼 나와있지만 웬만해서는 그렇게 하면 안된다. 빌어먹을 document

 

특히 pytorch model 내부에 조건부로 동작하는 모듈이 있을 경우 무조건 안된다.

 

예를 들어, 위 예시는 yolov8 모델인데 마지막에 NMS를 포함해서 ONNX로 변환한다. 이 때 만약 변환에 사용되는 예시 입력 이미지가 random noise이면 NMS에 들어가는 입력이 없기 때문에 변환 과정에서는 NMS가 동작하지 않는다.

 

따라서 변환 과정에서 쓸모없는 모듈이라고 여겨 삭제해버리는 것이다. 

 

이 경우, 입력에 검출 대상이 되는 객체가 많이 있는 이미지를 x 로 만들어서 torch.onnx.export 에 사용해야 한다. 

 

다시 말하면 변환에 사용되는 입력 x는 pytorch model 내부 모든 모듈을 한 번씩 다 통과하면서 연산될 이미지를 사용해야 한다. 

 

아쉽게도 x 자리에 여러개의 입력을 넣을 순 없기 때문에 만약 최고의 x를 찾기가 어렵다면 x를 바꿔가면서 계속해서 변환해보고, 가장 변환이 잘된 ONNX 모델을 사용하는 것이 맞다.