[ONNX] ONNX로 pytorch 딥러닝 모델 변환/추출/실행하기 (Yolov5)

 

다른 개발자 분과 이야기하다가 ONNX에 대해 접하게 되었습니다.

사실 ONNX를 코드에서 봤을 때는 뭐 새로운 게 딥러닝 라이브러리인가? 싶었는데, 생각보다 활용성이 많은 것 같아, 

이번에 한번 간단하게 테스트를 해보고자 합니다. 

 

간단하게 ONNX의 개념을 살펴보고, ONNX변환 코드를 소개하겠습니다. 

 

---

 

ONNX 란?

ONNX는  다른 환경에서 제작한 딥러닝 모델을 호환해서 사용할 수 도록 해주는 교차 플랫폼 모델 가속기입니다.

ONNX는 pytorch, tensorflow/keras, tflit, scikit-learn 및 기타 프레임워크 모델과 함께 사용할 수 있습니다. 

 

 

Pytorch에서 작성한 모델을 ONNX로 변환하는 방법

1. onnx 파일 추출 
2. onnxruntime 이용하여 실행

 

onnx 파일 추출 코드

onnx 파일 추출 코드입니다. 

yolov5 모델을 불러온 뒤, Conver_ONNX를 사용하여 변환이 가능합니다. 

 

import torch.onnx

def Convert_ONNX(model):

    # Let's create a dummy input tensor
    dummy_input = torch.randn(1, 3, 224,224, requires_grad=True, device='cpu')
    # Export the model
    torch.onnx.export(model,         # model being run
         dummy_input,       # model input (or a tuple for multiple inputs)
         "ImageClassifier.onnx",       # where to save the model
         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 = ['modelInput'],   # the model's input names
         output_names = ['modelOutput'], # the model's output names
         dynamic_axes={'modelInput' : {0 : 'batch_size'},    # variable length axes
                                'modelOutput' : {0 : 'batch_size'}})
    print(" ")
    print('Model has been converted to ONNX')


if __name__ == "__main__":

    # Let's load the model

    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', device='cpu')  # 예측 모델
    # Conversion to ONNx
    Convert_ONNX(model)

 

 

onnx 실행 및 테스트 코드

 

import onnxruntime
import numpy as np
import cv2
import torch
from time import time
ort_session = onnxruntime.InferenceSession("ImageClassifier.onnx" )


print(onnxruntime.get_device())
def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

# ONNX 런타임에서 계산된 결과값
x = torch.randn(1, 3, 224, 224, requires_grad=True) #

ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
pre_time = time()
ort_outs = ort_session.run(None, ort_inputs)
print("processing time : ", time()-pre_time)

yolo_model = torch.hub.load('ultralytics/yolov5', 'yolov5s', device='cpu' )  # 예측 모델
pre_time = time()
torch_out = yolo_model(x)
print("processing time : ", time()-pre_time)
np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-02, atol=1e-02)

print("Exported model has been tested with ONNXRuntime, and the result looks good!")

processing time :  0.007970809936523438 # onnx 모델 실행시간 processing time :  0.031972646713256836 # pytorch 모델 실행시간 Exported model has been tested with ONNXRuntime, and the result looks good!

 

테스트해보니, 실행시간에 차이가 꽤 났습니다.

가속기라는 말이 의미가 있어보입니다. 

또한, 예측 결과 값에 대해서도 유사하게 나타나는 듯합니다. 

 

현재 테스트한 환경은 cpu 환경입니다. 

gpu 환경에서 테스트해 보면 좋을 것 같습니다.