[ONNX] pytorch 모델을 ONNX로 변환하고 실행하기

ONNX (Open Neural Network eXchange)

 

ONNX는 기계 학습 모델을 표현하기 위해 만들어진 오픈 포맷으로 ONNX 런타임은 여러 다양한 플랫폼과 하드웨어(윈도우, 리눅스, 맥을 비롯한 플랫폼 뿐만 아니라 CPU, GPU 등의 하드웨어)에서 효율적인 추론을 가능하게 한다. 때문에 다양한 프레임워크와의 연계가 필요할 때 ONNX를 사용한다. (pytorch ↔︎ tensorflow ↔︎ caffe2 ↔︎ MXNet ↔︎ ...)

 

 

*참고

• https://pytorch.org/docs/stable/onnx.html
• https://tutorials.pytorch.kr/advanced/super_resolution_with_onnxruntime.html

 

 

ONNX 예제

 

1. Pytorch 모델 구현

 

# 필요한 import문
import io
import numpy as np

from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.onnx

import torch.nn as nn
import torch.nn.init as init

# 정의된 모델을 사용하여 초해상도 모델 생성
torch_model = SuperResolutionNet(upscale_factor=3)

 

 

2. 학습된 웨이트 불러오기

 

# 미리 학습된 가중치를 읽어옵니다
model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'
batch_size = 1    # 임의의 수

# 모델을 미리 학습된 가중치로 초기화합니다
map_location = lambda storage, loc: storage
if torch.cuda.is_available():
    map_location = None
torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))

# 모델을 추론 모드로 전환합니다
torch_model.eval()

 

 

3. 모델 변환

 

• Tracing이나 Scripting을 통해서 PyTorch 모델을 변환할 수 있는데 이 예제에서는 tracing을 통해 변환된 모델을 사용
• 모델을 변환하기 위해서는 torch.onnx.export() 함수를 호출
• 이 함수는 모델을 실행하여 그 실행을 추적(trace)한 다음 추적된 모델을 지정된 파일로 내보냄
• export 함수가 모델을 실행하기 때문에, 우리가 직접 텐서를 입력값으로 넘겨주어야 하고 이 텐서의 값은 알맞은 자료형과 shape이라면 랜덤한 값이어도 무관

# 모델에 대한 입력값
x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)
torch_out = torch_model(x)

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

저장된 onnx 파일에는 모델의 매개변수와 네트워크 구조를포함하는 바이너리 프로토콜 버퍼가 포함되어 있음

 

또한 layer 간의 입출력 크기를 확인하기 위해서 저장된 ONNX를 다시 불러와서 아래와 같은 방식으로 shape 정보를 저장하는 과정이 필요

from onnx import shape_inference

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

 

 

4. ONNX 모델 확인

 

• ONNX 런타임에서 변환된 모델을 사용했을 때 같은 결과를 얻는지 확인하기 위해서 torch_out 를 계산
• ONNX 런타임에서의 모델 결과값을 확인하기 전에 먼저 ONNX API를 사용해 ONNX 모델을 확인
  • onnx.load("super_resolution.onnx") 는 저장된 모델을 읽어온 후 머신러닝 모델을 취합하여 저장하고 있는 상위 파일 컨테이너인 onnx.ModelProto를 반환
  • onnx.checker.check_model(onnx_model) 는 모델의 구조를 확인하고 모델이 유효한 스키마(valid schema)를 가지고 있는지를 체크

 

import onnx

onnx_model = onnx.load("super_resolution.onnx")
onnx.checker.check_model(onnx_model)

 

 

5. ONNX 런타임과 pytorch 출력 비교

 

• ONNX 런타임의 Python API를 통해 결과값을 계산
• 이 부분은 보통 별도의 프로세스 또는 별도의 머신에서 실행되지만, 이 튜토리얼에서는 모델이 ONNX 런타임과 PyTorch에서 동일한 결과를 출력하는지를 확인하기 위해 동일한 프로세스에서 계속 실행
• 모델을 ONNX 런타임에서 실행하기 위해서는 미리 설정된 인자들로 모델을 위한 추론 세션을 생성
• 세션이 생성되면, 모델의 run() API를 사용하여 모델을 실행
• 이 API의 리턴값은 ONNX 런타임에서 연산된 모델의 결과값들을 포함하고 있는 리스트
• PyTorch와 ONNX 런타임에서 연산된 결과값이 서로 일치하는지 오차범위 (rtol=1e-03, atol=1e-05) 이내에서 확인 

import onnxruntime

ort_session = onnxruntime.InferenceSession("super_resolution.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(x)}
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!")

 

*주의 사항

PyTorch 모델은 NumPy 또는 Python 유형 및 함수를 사용하여 작성할 수 있지만 tracing 중에 NumPy 또는 Python 유형의 모든 변수(torch.Tensor가 아님)는 상수로 변환되므로 해당 값이 다음에 따라 변경되어야 하는 경우 잘못된 결과가 생성된다.