ONNX란?

먼저 ONNX에 대해 알아보겠습니다.

ONNX란 마이크로소프트와 페이스북이 공동으로 개발한 머신러닝 모델 공개 표준입니다. Tensorflow, PyTorch, SciKit-Learn, MXNet, Caffe2 등의 다양한 머신러닝 프레임워크를 지원하여 플랫폼과 디바이스 간의 상호운용이 가능하게 합니다. 우리는 각종 플랫폼(클라우드 or 엣지, CPU or GPU, etc) 상에서 최적의 성능을 필요로하지만 플랫폼마다 기능과 그 특성이 다르기 때문에 쉽지 않은 문제입니다. 따라서 ONNX는 이 문제를 해결하는데 많은 도움을 줄 수 있습니다.

ONNX를 지원하며 컨버팅할 수 있는 프레임워크들을 살펴봅시다.

컨버팅된 ONNX를 사용하여 추론할 수 있도록 하는 엔진들은 다음과 같습니다.

오늘은 ONNX를 사용하여 PyTorch 모델을 OpenCV에서 로드하고 추론엔진을 제작하는 방법에 대해 살펴보겠습니다.

PyTorch 모델의 저장과 불러오기

딥러닝 모델을 파일로 저장하는 방식으로는 학습된 파라미터만 저장하는 방식과 전체구조와 파라미터를 모두 저장하는 방식, 두 가지가 존재합니다.

학습 파라미터만 저장하는 경우

학습된 파라미터만 파일로 저장하고, 모델구조는 소스코드에 선언된 클래스(혹은 코드)를 이용하는 경우입니다.

# model.param 으로 저장
params = model.state_dict()
torch.save(params, "model.param")

# 불러오기
model = ModelClass(*args, **kwargs)
model.load_state_dict(torch.load("model.param"))
model.eval()

위와 같은 로직으로 일반적으로는 가능하지만, 현재 준비된 모델구조 dictionary key와 불어들이려는 파라미터 데이터 상의 dictionary key값이 불일치하는 부분이 하나라도 존재하는 경우 에러가 발생할 수 있습니다. 따라서 일일이 확인하여 key가 일치하는 경우에만 값을 세팅하도록 하여 안전하게 불어들일 수 있습니다. 해당 코드는 다음과 같습니다.

model = ModelClass(*args, **kwargs)
state_dict = torch.load(PATH, map_location = 'cpu')
model_dict = {} 
new_state_dict = net.state_dict() 
for (k,v) in checkpoint['state_dict'].items(): 
	print(k) 
    if k[7:] in state_dict: 
    	model_dict[k[7:]] = v 
new_state_dict.update(model_dict) 
model.load_state_dict(state_dict) 
model.eval()

모델구조 및 파라미터를 모두 저장하는 경우

# save
torch.save(model, "model.param")

# load
model = torch.load("model.param")

사용법이 간단하지만 딥러닝 모델구조를 정의하는 클래스 등이 PyTorch 업데이트를 통해 변경될 경우 정상적으로 실행되지 않을 가능성도 있으므로 파라미터를 개별적으로 저장하는 것을 더욱 권장합니다.

PyTorch 모델을 ONNX 형식으로 변환

PyTorch가 제공하는 Tracing이라는 방법을 통해 ONNX 형식으로 변환이 가능합니다. 모델을 변환하기 위해서는 torch.onnx.export() 함수를 호출합니다. PyTorch는 동적계산그래프 방식을 사용하므로 모델을 실행하여 어떤 연산자들이 출력값을 계산하는데 사용되는지를 기록하며 ONNX 포맷을 생성합니다. export 함수가 모델을 실행하므로 파라미터로 입력 텐서를 넘겨주어야 하고, 이 텐서는 실제 이미지가 아닌 랜덤한 값이어도 무방합니다. ONNX로 변환된 그래프의 경우 export 함수에 전달한 입력값의 크기로 고정된다는 것을 기억해야 합니다.

import torch.onnx

dummy = torch.empty(1, 3, 224, 224, dtype = torch.float32)
torch.onnx.export(model, dummy, "model.onnx")

입력데이터와 출력데이터의 이름을 특정 이름으로 지정하고 싶다면 다음과 같이 호출하는 것이 가능합니다.

# 입출력이 1개인 경우
torch.onnx.export(model, dummy,  input_names = ['input'], output_names = ['output'], "model.onnx")

# 입력이 1개 출력이 2개인 경우
torch.onnx.export(model, dummy,  input_names = ['input'], output_names = ['cls_score','bbox_pred'], "model.onnx")

변환된 ONNX 파일의 검증

단순히 파일이 생성된 것만으로 ONNX가 정상적으로 변환되었는지, 추론 성능의 변화가 없는지 확인할 수 없습니다. 파이토치에서는 변환된 모델의 검증방법을 제공하고 있는데, 기존 모델과 변환된 ONNX 모델의 결과값을 비교하여 값의 차이가 오차허용 범위 이내라면 정상적으로 변환된 것으로 판단하도록 하고 있습니다. 아래는 관련 코드입니다.

ort_session = onnxruntime.InferenceSession("feathernets.onnx") 
 def to_numpy(tensor): 
   return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy() 
 # ONNX 런타임에서 계산된 결과값 ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(dummy_input)} 
ort_outs = ort_session.run(None, ort_inputs) 
# ONNX 런타임과 PyTorch에서 연산된 결과값 비교 
np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05) 
print("Exported model has been tested with ONNXRuntime, and the result looks good!")

오차범위에서 벗어난 경우 위의 코드에서 np.testing.assert_allclose 함수가 오류를 발생시키므로 맨 아래의 문장이 출력된다면 정상적인 변환으로 판단합니다. 이번 작업에서도 정상으로 변환되었음을 확인할 수 있었습니다.

ONNX에 Shape 정보 저장

위와 같은 과정으로 ONNX 파일을 얻을 수 있었습니다. 하지만 이 경우 layer간 입출력 크기를 확인할 수 없어서 확인할 수 있는 세부정보가 제한적입니다.

모델 전체 구조를 확인하기 위해서는 layer간 입출력 크기 정보가 중요합니다. 해당 정보 저장을 위해서는 저장된 ONNX를 다시 불러들인 후, 아래와 같은 방식으로 새롭게 저장합니다.

import onnx
from onnx import shape_inference
onnx.save(onnx.shape_inference.infer_shapes(onnx.load("model.onnx")), "model.onnx")

ONNX 파일 확인

shell에서 확인

생성된 ONNX 파일 내부 정보를 확인하도록 합니다.
onnx 패키지를 설치합니다.

pip install onnx

onnx를 불러들인 후, numpy_helper를 통해 각 layer 값을 numpy 자료 형으로 변환 후 출력합니다.

onnx_model = onnx.load("model.onnx")
graph = onnx_model.graph
initializers = dict()
for init in graph.initializer:
    initializers[init.name] = numpy_helper.to_array(init)
print(initializers.keys())

GUI로 확인

netron을 설치합니다.

snap install netron

netron을 사용하여 ONNX 파일을 실행하면 모델을 GUI 형태로 확인할 수 있습니다.

netron model.onnx

OpenCV를 통해 불러오기

#include <iostream>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/dnn/dnn.hpp>

using namespace std;
using namespace cv;
using namespace dnn;

int main()
{
    // load the neural network model
    cv::dnn::Net net = cv::dnn::readNetFromONNX("model.onnx");
    net.setPreferableBackend(dnn::DNN_BACKEND_OPENCV);
    net.setPreferableTarget(dnn::DNN_TARGET_CPU);
    
    Mat input;
    input = dnn::blobFromImage(src); // src는 cv::Mat 이미지 데이터
    this->Lenet5.setInput(input);
	Mat output = net.forward();
}

필요한 경우 아래와 같이 학습시 사용한 이미지에 적용된 것과 동일한 정규화를 수행해야 합니다. 여기서의 정규화는 0~255 값을 0과 1사이로 scaling하는 것이고, 표준화는 데이터 값에서 평균을 뺀 후, 표준편차로 나누는 작업입니다. 이런 전처리 작업은 모델마다 다르므로 학습시 수행한 전처리를 참고하여 동일하게 적용해 주는 것이 중요합니다.

cv::Mat blob = cv::dnn::blobFromImage(face, 1.0 / 255, cv::Size(224, 224), cv::Scalar(0, 0, 0), true, false, CV_32F);  
for (int r = 0; r < 224; r++) { 
    for (int c = 0; c < 224; c++) { 
	blob.at<float>(Vec<int, 4>(0, 0, r, c)) = (blob.at<float>(Vec<int, 4>(0, 0, r, c)) - 0.14300402) / 0.10050353; 
	blob.at<float>(Vec<int, 4>(0, 1, r, c)) = (blob.at<float>(Vec<int, 4>(0, 1, r, c)) - 0.1434545) / 0.100842826; 
	blob.at<float>(Vec<int, 4>(0, 2, r, c)) = (blob.at<float>(Vec<int, 4>(0, 2, r, c)) - 0.14277956) / 0.10034215; 
    } 
}

Reference

• https://tutorials.pytorch.kr/advanced/super_resolution_with_onnxruntime.html
• https://discuss.pytorch.org/t/solved-keyerror-unexpected-key-module-encoder-embedding-weight-in-state-dict/1686/3
• https://gaussian37.github.io/dl-pytorch-deploy/