728x90
반응형
다른 개발자 분과 이야기하다가 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 환경에서 테스트해 보면 좋을 것 같습니다.
728x90
반응형
'코드 > ONNX' 카테고리의 다른 글
[ONNX] ONNX Runtime? Pytorch 모델 ONNX 변환하기 (Renset18) (0) | 2023.08.08 |
---|---|
[ONNX] ONNX 란? - 탄생 배경, 목적, 성능 테스트 (0) | 2022.12.16 |
댓글