[ONNX] 딥러닝 모델 ONNX Runtime으로 CPU 환경에서 가속화하기

요즘은 어지간한 딥러닝 모델을 GPU 없이 돌리기 어렵지만, 또 의외로 가벼운 모델들은 CPU 만으로 돌릴 수 있다. 가능하다면 클라우드 비용도 줄일 수 있으니 온라인 예측이 필요한 경우가 아니라면 CPU 환경에서 인퍼런스하는 것도 고려해 볼 만하다.

 

물론 CPU로 딥러닝 모델 인퍼런스를 하게 되면 상당히 느리다. 때문에 ONNX 모델 변환을 하고, ONNX runtime으로 인퍼런스를 수행하면 조금이라도 모델 인퍼런스 속도를 향상시킬 수 있다. 또한 TensorRT와 달리 ONNX 모델 변환의 경우 입력 텐서 크기 또한 동적으로 가져갈 수 있다는 장점이 있다.

 

물론 하드웨어 환경에 따라, 모델에 따라, 입력 텐서의 크기에 따라 속도 향상의 정도가 다르거나, 오히려 속도가 느려질 수도 있으니 테스트를 해봐야 한다.

 

Resnet 으로 간단하게 테스트해봤을 때 약 1.5~1.7배 정도의 속도 향상이 있었고, 현재 사용중인 CNN 기반의 검출기로 테스트를 해봤을 때도 비슷한 정도로 속도가 향상되었다.

 

속도가 많이 빠를 필요 없고, 모델이 어느정도 가볍다면 CPU 환경에서 ONNX 런타임으로 모델을 배포하는 것도 충분히 생각해볼 수 있는 옵션인 것 같다.

 

ONNX Runtime 예제 코드

import torch
import torchvision
import numpy as np
import onnx
import onnxruntime as ort
from onnx import shape_inference
import time


# PyTorch 모델 로드
torch_model = torchvision.models.resnet18(pretrained=False)
torch_model.eval()

# 예제 입력 데이터 생성
dummy_input = torch.randn(1, 3, 500, 500, requires_grad=True)

repetitions = 10

for _ in range(5):
    _ = torch_model(dummy_input)

start = time.time()
with torch.no_grad():
    for rep in range(repetitions):
        torch_out = torch_model(dummy_input)
end = time.time()

print('torch 모델 평균 소요 시간 : ', (end-start)/repetitions)
    

# # 모델 변환
torch.onnx.export(torch_model,               # 실행될 모델
                    dummy_input,                         # 모델 입력값 (튜플 또는 여러 입력값들도 가능)
                    "test_resnet18.onnx",   # 모델 저장 경로 (파일 또는 파일과 유사한 객체 모두 가능)
                    export_params=True,        # 모델 파일 안에 학습된 모델 가중치를 저장할지의 여부
                    opset_version=10,          # 모델을 변환할 때 사용할 ONNX 버전
                    do_constant_folding=True,  # 최적화시 상수폴딩을 사용할지의 여부
                    input_names = ['input'],   # 모델의 입력값을 가리키는 이름
                    output_names = ['output'], # 모델의 출력값을 가리키는 이름
                    dynamic_axes={'input' : {0: 'batch_size', 2: 'height', 3: 'width'}},    # 가변적인 길이를 가진 차원
                    )


path = "./test_resnet18.onnx"
onnx.save(onnx.shape_inference.infer_shapes(onnx.load(path)), path)

# # ONNX 모델 로드
onnx_model = onnx.load("./test_resnet18.onnx")
onnx.checker.check_model(onnx_model)

# ONNX 런타임 세션 열기 (CPU 사용 설정)
ort_session = ort.InferenceSession("./test_resnet18.onnx", providers=['CPUExecutionProvider'])
print(ort.get_device())

# 인퍼런스 실행
ort_inputs = {ort_session.get_inputs()[0].name: np.array(dummy_input.detach())}

for _ in range(5):
    _ = ort_session.run(None, ort_inputs)

start = time.time()
with torch.no_grad():
    for rep in range(repetitions):
        ort_outputs = ort_session.run(None, ort_inputs)
end = time.time()

print('ONNX 평균 소요 시간 : ', (end-start)/repetitions)