PyTorch 모델을 ONNX로 바꿔 FastAPI 추론 2~3배 빠르게 (변환 코드 + 에러 잡기)

by

PyTorch 모델을 ONNX로 바꿔 FastAPI 추론 2~3배 빠르게 (변환 코드 + 에러 잡기)

FastAPI로 PyTorch 모델을 서빙하다 보면 추론 지연 시간 때문에 곤란해질 때가 많습니다. 특히 CPU 서버에서는 torch.no_grad()만으로 버티기 힘든 경우가 있는데, 이럴 때 PyTorch 모델을 ONNX로 변환해서 ONNX Runtime으로 추론하면 별도의 모델 재학습 없이도 체감 속도가 눈에 띄게 좋아집니다. 이 글에서는 실제로 동작하는 변환 코드, FastAPI 연동 방식, 그리고 변환 과정에서 자주 만나는 에러와 해결법을 정리합니다.

왜 ONNX 변환만으로 속도가 오르는가

PyTorch의 즉시 실행(eager execution) 방식은 파이썬 인터프리터가 레이어마다 연산 그래프를 매번 다시 해석합니다. 반면 ONNX(Open Neural Network Exchange)는 모델을 정적 계산 그래프로 직렬화한 표준 포맷이고, ONNX Runtime은 이 그래프를 미리 최적화(연산자 퓨전, 상수 폴딩, 불필요한 노드 제거)한 뒤 실행합니다. 파이썬 오버헤드가 줄고 CPU에서는 MLAS, GPU에서는 CUDA/TensorRT 실행 provider가 커널을 선택하기 때문에, 같은 하드웨어에서도 배치 크기가 작은 실시간 추론 시나리오일수록 체감 개선 폭이 큽니다. ONNX 포맷 자체의 명세는 위키백과의 Open Neural Network Exchange 문서에 정리되어 있습니다.

주의할 점은 “무조건 2~3배”가 보장되는 건 아니라는 것입니다. 커스텀 연산자가 많거나 이미 GPU에서 배치 처리가 큰 모델이라면 효과가 작을 수 있고, 반대로 CPU 단건 추론처럼 오버헤드 비중이 큰 구조일수록 개선 폭이 커집니다.

PyTorch 모델을 ONNX로 변환하는 실전 코드

먼저 필요한 패키지를 설치합니다.

pip install onnx onnxruntime
# GPU 추론이 필요하면
pip install onnxruntime-gpu

변환은 torch.onnx.export로 진행합니다. 배치 크기가 요청마다 달라질 수 있으므로 dynamic_axes를 반드시 지정해야 합니다.

import torch

model.eval()  # 반드시 eval 모드로 전환 (BatchNorm/Dropout 동작 고정)

dummy_input = torch.randn(1, 3, 224, 224)  # 모델 입력 shape에 맞게 조정

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"},
    },
    opset_version=17,
    do_constant_folding=True,
)

파이썬 딥러닝 코드가 실행되는 터미널 화면

변환된 파일은 onnx.checker로 그래프 유효성을 검증하는 습관을 들이는 게 좋습니다.

import onnx

onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX 모델 검증 통과")

FastAPI에서 ONNX Runtime으로 추론하기

핵심은 InferenceSession을 요청마다 생성하지 않고 앱 시작 시 한 번만 로드하는 것입니다. FastAPI의 lifespan을 활용하면 프로세스 생명주기 동안 세션을 재사용할 수 있습니다.

from contextlib import asynccontextmanager
import numpy as np
import onnxruntime as ort
from fastapi import FastAPI
from pydantic import BaseModel

session = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global session
    session = ort.InferenceSession(
        "model.onnx",
        providers=["CPUExecutionProvider"],  # GPU는 ["CUDAExecutionProvider", "CPUExecutionProvider"]
    )
    yield
    session = None

app = FastAPI(lifespan=lifespan)

class InferenceRequest(BaseModel):
    data: list

@app.post("/predict")
def predict(req: InferenceRequest):
    input_array = np.array(req.data, dtype=np.float32)  # dtype 명시가 핵심
    outputs = session.run(None, {"input": input_array})
    return {"result": outputs[0].tolist()}

여기서 session.run(None, ...)의 첫 인자를 None으로 두면 모든 출력 노드를 반환합니다. 특정 출력만 필요하면 ["output"]처럼 이름 리스트로 지정해 불필요한 텐서 복사를 줄일 수 있습니다. ONNX Runtime의 실행 provider와 세션 옵션에 대한 상세 설명은 ONNX Runtime 공식 문서를 참고하면 됩니다.

변환·추론 중 자주 만나는 에러와 해결법

1) RuntimeError: Exporting the operator ... to ONNX opset version X is not supported
모델이 사용하는 연산자가 지정한 opset_version에 없을 때 발생합니다. opset_version을 17 이상으로 올리거나, torch와 onnx 패키지를 최신 버전으로 업그레이드한 뒤 다시 시도하면 대부분 해결됩니다. PyTorch가 지원하는 opset별 연산자 목록은 PyTorch 공식 ONNX 문서에서 확인할 수 있습니다.

2) onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: Got invalid dimensions for input
변환 시 배치 축을 고정한 채(dynamic_axes 미지정) export했는데 실제 요청의 배치 크기가 다를 때 발생합니다. 위 코드처럼 dynamic_axes에 배치 차원을 "batch_size"로 명시하고 다시 export해야 합니다.

3) InvalidArgument: Unexpected input data type. Actual: (tensor(double)) , expected: (tensor(float))
numpy 배열을 만들 때 dtype을 지정하지 않으면 기본값이 float64가 되는데, 대부분의 PyTorch 모델은 float32 가중치로 학습됩니다. np.array(data, dtype=np.float32)처럼 dtype을 명시하면 해결됩니다.

서버 API 성능 최적화 벤치마크 그래프

4) TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect
모델의 forward 안에 텐서 값에 따라 분기하는 if문이나 파이썬 for 루프가 있을 때 발생합니다. torch.onnx.export는 트레이싱 방식이라 한 번 실행된 경로만 그래프에 고정되므로, 조건 분기가 실제로 있는 모델이라면 torch.jit.script로 스크립팅한 뒤 export하거나, 분기를 torch.where 같은 텐서 연산으로 바꿔야 안전합니다.

5) CUDAExecutionProvider 지정했는데도 CPU로 동작
onnxruntime(CPU 전용 패키지)이 설치된 상태에서 providers=["CUDAExecutionProvider"]를 넘기면 조용히 CPU로 폴백됩니다. pip uninstall onnxruntime 후 onnxruntime-gpu를 설치하고, ort.get_available_providers()로 실제 사용 가능한 provider 목록을 먼저 확인하는 것이 확실합니다.

방식 비교와 실제 적용 팁

항목 PyTorch eager 추론 ONNX Runtime 추론
요청당 파이썬 오버헤드 레이어마다 발생 그래프 실행 1회로 최소화
그래프 최적화 없음(TorchScript 별도 적용 필요) 로드 시 연산자 퓨전·상수 폴딩 자동 적용
동적 입력 크기 자유로움 dynamic_axes 사전 지정 필요
커스텀/희귀 연산자 제약 없음 opset 미지원 시 변환 실패 가능
배포 파일 .pt/.pth + 모델 클래스 코드 단일 .onnx 파일

CPU 서버에서 배치 크기 1~4의 짧은 요청을 처리하는 API라면 ONNX Runtime 전환 효과가 가장 크게 나타납니다. 반대로 이미 배치를 크게 묶어 GPU에서 처리하는 파이프라인이라면 변환 전에 실제 요청 패턴으로 벤치마크를 먼저 돌려보고 결정하는 것이 안전합니다. 벤치마크는 time.perf_counter()로 100회 이상 반복 호출해 평균과 p95를 함께 재는 방식을 권장합니다.

마무리: 바로 적용할 다음 단계

핵심은 세 가지입니다. model.eval() 상태에서 dynamic_axes를 지정해 export하고, onnx.checker로 그래프를 검증하고, FastAPI에서는 InferenceSession을 앱 시작 시 한 번만 만들어 재사용하는 것입니다. 오늘 바로 적용하려면 기존 추론 엔드포인트를 복제해 /predict-onnx로 새 라우트를 하나 추가하고, 두 엔드포인트의 응답 시간을 동일한 입력으로 비교해보는 것부터 시작하면 됩니다. 결과값(logits/확률)이 PyTorch 버전과 오차 범위(np.allclose 기준 1e-4 정도) 안에서 일치하는지도 함께 확인해야 변환 과정에서 정밀도 손실이 없었는지 검증할 수 있습니다.

자주 묻는 질문(FAQ)

Q1. ONNX로 변환하면 모델 정확도가 떨어지나요?
opset이 지원하는 연산자로만 이루어진 모델이라면 대부분 수치적으로 거의 동일한 결과를 냅니다.다만 변환 후에는 항상 원본 PyTorch 출력과 np.allclose로 비교 검증하는 것이 안전합니다.

Q2. 모든 PyTorch 모델을 ONNX로 변환할 수 있나요?
아닙니다. 커스텀 CUDA 커널이나 파이썬 제어 흐름에 강하게 의존하는 모델, 최신 opset에 아직 없는 연산자를 쓰는 모델은 변환이 실패하거나 트레이싱 결과가 부정확할 수 있습니다.

Q3. GPU 서버에서도 ONNX 변환이 의미가 있나요?
있습니다. CUDAExecutionProvider나 TensorRT provider를 쓰면 그래프 최적화와 커널 선택 덕분에 순수 PyTorch 대비 지연 시간이 줄어드는 경우가 많습니다. 다만 배치가 이미 큰 처리량 중심 워크로드라면 개선 폭이 CPU 환경보다는 작을 수 있습니다.

Leave a Comment