PyTorch 모델을 ONNX로 변환하고 ONNX Runtime으로 배포하는 방법

딥러닝 모델을 연구하고 개발할 때 PyTorch는 강력하고 유연한 환경을 제공합니다. 하지만 실제 서비스 환경에 모델을 배포하거나, 모바일 및 임베디드 기기처럼 자원이 제한된 환경에서 모델을 구동할 때는 프레임워크의 무거운 무게가 걸림돌이 되기도 합니다. 이러한 이종 환경 간의 장벽을 허물고 배포 효율성을 극대화하기 위해 개발된 표준이 바로 ONNX입니다. 이번 글에서는 PyTorch 모델을 ONNX 포맷으로 변환하고, ONNX Runtime을 통해 고성능 인퍼런스를 구현하는 방법을 구체적으로 알아보겠습니다.

Generated by Gemini AI.

핵심 요약 3줄

• ONNX는 서로 다른 딥러닝 프레임워크 간에 모델을 공유하고 실행할 수 있도록 지원하는 개방형 표준 포맷입니다.
• PyTorch의 torch.onnx.export 함수를 사용하면 몇 줄의 코드만으로 기존 모델을 ONNX 파일로 내보낼 수 있습니다.
• 변환된 모델은 ONNX Runtime 엔진을 활용하여 CPU, GPU 등 다양한 하드웨어 환경에서 최적화된 속도로 실행됩니다.

1. ONNX란?

ONNX(Open Neural Network Exchange)는 다양한 딥러닝 프레임워크 간의 모델 호환성을 제공하는 개방형 표준 포맷입니다. Microsoft와 Facebook이 공동으로 개발한 이 표준은 PyTorch, TensorFlow, MXNet, Caffe2 등 여러 프레임워크에서 모델을 변환하고 배포할 수 있도록 지원합니다.

ONNX를 활용하면 특정 프레임워크에서 학습한 모델을 다른 환경에서 실행할 수 있으며, 특히 경량화된 환경이나 고성능 최적화가 필요한 배포 환경에서 유용합니다. 또한, ONNX Runtime과 같은 고성능 엔진을 통해 다양한 플랫폼에서 효율적으로 모델을 실행할 수 있습니다.

2. ONNX의 주요 특징 및 프레임워크 비교

ONNX는 모델 배포 파이프라인을 단순화하는 여러 특징을 가지고 있습니다. 기존 프레임워크들과의 차이점과 주요 특징을 표로 정리했습니다.

구분	PyTorch / TensorFlow	ONNX
주요 목적	모델 연구, 설계, 학습(Training)	모델 상호 운용, 최적화, 배포(Inference)
독립성	해당 프레임워크 전용 환경 필요	프레임워크 독립적 (다양한 환경 지원)
실행 엔진	프레임워크 내장 백엔드	ONNX Runtime 및 하드웨어 가속기 파트너십
최적화 기법	그래프 최적화 및 퀀타이제이션 지원	상수 폴딩, 노드 퓨전 등 하드웨어 맞춤형 최적화 내장

3. PyTorch 모델을 ONNX로 변환하는 방법

PyTorch에서 학습한 모델을 ONNX 형식으로 변환하는 과정은 비교적 간단합니다. torch.onnx.export 함수를 사용하여 변환할 수 있으며, 변환된 ONNX 모델은 ONNX Runtime을 활용하여 실행할 수 있습니다.

3.1. 변환을 위한 사전 준비

먼저 모델 변환과 검증, 실행에 필요한 핵심 라이브러리들을 설치해야 합니다. 터미널에서 아래 명령어를 실행합니다.

Bash

pip install onnx onnxruntime torch torchvision

3.2. PyTorch 모델 정의

간단한 선형 레이어와 활성화 함수로 구성된 PyTorch 모델을 정의합니다. 모델 변환 전에는 반드시 평가 모드(model.eval())로 전환해야 합니다. 드롭아웃이나 배치 정규화 같은 레이어들이 평가 시점의 동작으로 고정되어야 정확한 그래프가 생성되기 때문입니다.

Python

import torch
import torch.nn as nn
import torch.onnx

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(3, 3)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(3, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 모델 인스턴스 생성 및 평가 모드 전환
model = SimpleModel()
model.eval()

3.3. PyTorch 모델을 ONNX로 변환

PyTorch는 모델의 연산 흐름을 추적하는 추적(Tracing) 방식을 사용합니다. 따라서 모델에 흘려보낼 가짜 데이터인 더미 입력 텐서가 필요합니다.

Python

# 더미 입력 텐서 생성 (배치 크기 1, 입력 차원 3)
dummy_input = torch.randn(1, 3)

# ONNX 변환 수행
torch.onnx.export(
    model,                       # 변환할 PyTorch 모델
    dummy_input,                 # 모델 입력 텐서 (더미 데이터)
    "simple_model.onnx",         # 변환된 모델 저장 파일명
    export_params=True,          # 모델의 학습된 가중치 저장 여부
    opset_version=11,            # ONNX Opset 버전 지정
    do_constant_folding=True,    # 상수 폴딩 최적화 옵션 활성화
    input_names=['input'],       # 입력 노드 이름 지정
    output_names=['output'],     # 출력 노드 이름 지정
    dynamic_axes={               # 동적 배치 크기 지원을 위한 설정
        'input': {0: 'batch_size'}, 
        'output': {0: 'batch_size'}
    }
)

위 코드를 실행하면 프로젝트 디렉토리에 simple_model.onnx 파일이 생성됩니다. dynamic_axes 옵션을 지정했기 때문에, 배포 시점에 다양한 배치 크기의 데이터를 유연하게 처리할 수 있습니다.

4. ONNX 모델 실행 및 검증

변환이 끝났다면 모델 파일이 깨지지 않고 정상적으로 구조를 갖추었는지 검증해야 합니다.

Python

import onnx

# ONNX 모델 로드
onnx_model = onnx.load("simple_model.onnx")

# 모델 구조 및 스키마 검증
onnx.checker.check_model(onnx_model)
print("ONNX 모델이 정상적으로 변환되었습니다.")

4.1. ONNX Runtime을 활용한 실행

검증을 마친 모델은 ONNX Runtime을 사용하여 실제로 추론을 수행할 수 있습니다. 이때 입력 데이터는 PyTorch 텐서가 아니라 넘파이(NumPy) 배열을 사용합니다.

Python

import onnxruntime as ort
import numpy as np

# ONNX Runtime 세션 생성
ort_session = ort.InferenceSession("simple_model.onnx")

# 입력 데이터 생성 (넘파이 배열 활용)
input_data = np.random.randn(1, 3).astype(np.float32)

# 모델 실행 및 결과 출력
outputs = ort_session.run(None, {"input": input_data})
print("ONNX 모델 출력:", outputs)

개발을 위한 팁

• Netron 시각화 도구 활용하기: 변환된 ONNX 모델의 내부 구조와 노드 연결 상태를 눈으로 확인하고 싶다면 Netron(https://netron.app/) 웹사이트를 이용해 보세요. 모델 파일을 드래그앤드롭하는 것만으로 입력 노드 명칭과 연산 흐름을 직관적으로 파악할 수 있습니다.
• Opset 버전 선택 기준: Opset 버지는 일종의 연산자 명세서 버전입니다. 무조건 최신 버전을 고르기보다는 배포 타깃 환경(예: 텐서RT, 모바일 SDK 등)이 공식적으로 지원하는 Opset 버전을 사전에 확인하고 맞추는 것이 호환성 측면에서 안전합니다.

흔히 하는 실수

• model.eval() 누락: 학습 모드 상태에서 변환을 진행하면 드롭아웃 등의 확률적 연산이 그래프에 그대로 포함되어 버립니다. 이로 인해 인퍼런스 환경에서 매번 다른 예측 값이 나오는 치명적인 문제가 발생할 수 있으니 변환 전 상태 전환을 명확히 해야 합니다.
• 동적 축(Dynamic Axes) 미지정: 변환할 때 dynamic_axes를 설정하지 않으면 더미 데이터로 넣었던 모양(Shape) 그대로 고정됩니다. 예를 들어 가짜 데이터를 배치 크기 1로 넣었다면, 실제 서비스에서 배치 크기 4로 요청이 들어왔을 때 형상 불일치 에러를 뱉으며 멈추게 됩니다.

5. 결론

ONNX는 PyTorch로 개발한 모델을 프레임워크 제약 없이 다양한 환경에 배포할 수 있도록 돕는 가교 역할을 합니다. torch.onnx.export 함수를 이용하면 기존 파이썬 코드를 효율적인 인퍼런스용 그래프로 쉽게 전환할 수 있고, ONNX Runtime 엔진과 결합하여 서버나 엣지 디바이스에서 한층 더 빠른 속도로 실행이 가능합니다.

프레임워크의 무거운 의존성에서 벗어나 실제 서비스 환경에 최적화된 형태로 모델을 패키징하고 배포하고 싶다면, 오늘 다룬 ONNX 변환 파이프라인을 프로젝트에 적극적으로 적용해 보시길 권장합니다.