Published on

PyTorch 모델을 ONNX+INT8로 4배 경량화하는 법

Authors

서빙 환경에서 PyTorch 모델을 그대로 올리면 모델 파일이 크고(배포/로딩 느림), FP32 연산 때문에 추론 비용이 커지는 경우가 많습니다. 특히 CPU 추론이나 엣지 배포에서는 메모리 압박이 곧바로 지연시간과 비용으로 이어집니다.

이 글에서는 PyTorch 모델을 ONNX로 내보낸 뒤 INT8 양자화를 적용해 모델 크기를 이론적으로 약 4배(32bit float → 8bit int) 줄이는 과정을 실전 관점에서 정리합니다. 또한 “파일 크기만 줄고 속도는 안 빨라진다” 같은 흔한 함정과 검증 방법까지 함께 다룹니다.

관련해서 TensorRT INT8에서 자주 터지는 이슈들은 아래 글이 보완재가 됩니다.


왜 ONNX+INT8인가

ONNX의 가치

  • 프레임워크 독립 포맷이라 런타임 선택지가 넓습니다(ONNX Runtime, TensorRT, OpenVINO 등).
  • PyTorch eager 모드의 오버헤드가 사라지고, 그래프 최적화(상수 폴딩, 연산 fuse 등)를 받을 수 있습니다.

INT8 양자화의 가치

  • 파라미터가 FP32에서 INT8로 바뀌면 가중치 저장 공간이 1/4이 됩니다.
  • CPU에서는 INT8 커널(예: AVX2/VNNI) 활용 시 지연시간이 크게 줄 수 있습니다.
  • GPU에서는 TensorRT 같은 엔진에서 INT8 최적화가 강력합니다(단, 캘리브레이션과 레이어 지원 여부가 중요).

주의할 점은, “모델 파일 크기 4배 감소”는 비교적 쉽게 달성되지만, “지연시간 4배 개선”은 하드웨어/연산자/배치 크기/런타임 최적화에 따라 편차가 큽니다.


전체 파이프라인 개요

  1. PyTorch 모델을 eval()로 전환하고 입력/출력 시그니처를 확정
  2. ONNX로 export (torch.onnx.export)
  3. ONNX 모델 검증(형상, 연산자, 정확도 스모크 테스트)
  4. 양자화 전략 선택
    • CPU 중심: ONNX Runtime quantize_dynamic 또는 quantize_static
    • GPU/TensorRT: INT8 엔진 빌드(캘리브레이션 필요)
  5. 성능/정확도 측정 및 회귀 방지(벤치마크 자동화)

이 글은 ONNX Runtime 기반 INT8 양자화(특히 CPU) 를 중심으로 설명하고, 마지막에 TensorRT 방향도 연결합니다.


1) PyTorch 모델을 ONNX로 내보내기

체크리스트

  • model.eval() 필수(드롭아웃/배치정규화 동작 고정)
  • 입력 텐서 dtype/shape 확정
  • 동적 배치가 필요하면 dynamic_axes 지정
  • export 후 onnx.checker로 검증

예시 코드: ONNX export

import torch
import onnx

# 예시 모델 (실전에서는 학습된 모델 로드)
class MLP(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(768, 512),
            torch.nn.ReLU(),
            torch.nn.Linear(512, 10),
        )

    def forward(self, x):
        return self.net(x)

model = MLP().eval()

dummy = torch.randn(1, 768)  # 배치 1 예시
onnx_path = "mlp_fp32.onnx"

torch.onnx.export(
    model,
    dummy,
    onnx_path,
    input_names=["input"],
    output_names=["logits"],
    opset_version=17,
    do_constant_folding=True,
    dynamic_axes={
        "input": {0: "batch"},
        "logits": {0: "batch"},
    },
)

m = onnx.load(onnx_path)
onnx.checker.check_model(m)
print("ONNX export OK:", onnx_path)

opset은 몇을 써야 하나

  • 일반적으로 opset_version은 17 이상을 권장합니다(런타임/도구 호환성 확인 필요).
  • 특정 연산자가 낮은 opset에서만 안정적인 경우도 있어, 배포 런타임 기준으로 고정하는 게 좋습니다.

2) ONNX Runtime로 FP32 추론 스모크 테스트

양자화 전에 “ONNX로 바꾼 것만으로 결과가 동일한지”부터 확인해야 합니다. 이 단계가 없으면, 이후 정확도 하락이 export 문제인지 양자화 문제인지 분리가 안 됩니다.

import numpy as np
import onnxruntime as ort

sess = ort.InferenceSession(
    "mlp_fp32.onnx",
    providers=["CPUExecutionProvider"],
)

x = np.random.randn(4, 768).astype(np.float32)
logits = sess.run(["logits"], {"input": x})[0]
print(logits.shape, logits.dtype)

여기서부터 이미 느리다면(혹은 모델이 너무 크다면) ONNX graph 최적화, provider 설정, 스레드 설정이 필요합니다.


3) INT8 양자화 전략 선택: Dynamic vs Static

Dynamic quantization

  • 가중치를 INT8로 바꾸고, 활성화(activation)는 런타임에서 동적으로 스케일을 잡습니다.
  • 장점: 캘리브레이션 데이터셋이 없어도 적용 가능, 적용이 쉬움
  • 단점: 정적 대비 성능/정확도 최적이 아닐 수 있음
  • 주로 MatMul, Gemm 계열(FC/Transformer의 일부)에 효과적

Static quantization

  • 가중치뿐 아니라 activation도 INT8로 내리기 위해 캘리브레이션을 수행합니다.
  • 장점: 성능 최적화 여지가 큼
  • 단점: 대표성 있는 캘리브레이션 데이터가 필요, 파이프라인이 복잡

현업에서는 “일단 dynamic으로 1차 경량화 및 성능 개선”을 만든 뒤, 더 필요하면 static 또는 TensorRT INT8로 넘어가는 흐름이 많습니다.


4) ONNX Runtime Dynamic INT8 양자화로 모델 크기 4배 줄이기

가장 빠르게 결과를 얻는 방법입니다.

from onnxruntime.quantization import quantize_dynamic, QuantType

fp32_path = "mlp_fp32.onnx"
int8_path = "mlp_int8_dynamic.onnx"

quantize_dynamic(
    model_input=fp32_path,
    model_output=int8_path,
    weight_type=QuantType.QInt8,  # 또는 QuantType.QUInt8
)

print("Saved:", int8_path)

파일 크기 비교 스크립트

import os

def size_mb(p):
    return os.path.getsize(p) / (1024 * 1024)

for p in ["mlp_fp32.onnx", "mlp_int8_dynamic.onnx"]:
    print(p, f"{size_mb(p):.2f} MB")

대부분의 dense 모델에서 가중치 비중이 크면 클수록 4배에 가까운 감소가 관찰됩니다. 다만 임베딩 테이블, 레이어 구성, 양자화되지 않는 노드 비율에 따라 감소폭은 달라집니다.


5) INT8 모델 정확도 검증: 수치 비교와 태스크 지표

양자화는 근사이므로 출력이 완전히 같을 수 없습니다. 따라서 “허용 가능한 오차 범위”를 정의하고 자동화하는 게 중요합니다.

로짓 수준 비교(간단 스모크)

import numpy as np
import onnxruntime as ort

x = np.random.randn(32, 768).astype(np.float32)

sess_fp32 = ort.InferenceSession("mlp_fp32.onnx", providers=["CPUExecutionProvider"])
sess_int8 = ort.InferenceSession("mlp_int8_dynamic.onnx", providers=["CPUExecutionProvider"])

y_fp32 = sess_fp32.run(["logits"], {"input": x})[0]
y_int8 = sess_int8.run(["logits"], {"input": x})[0]

mae = np.mean(np.abs(y_fp32 - y_int8))
maxe = np.max(np.abs(y_fp32 - y_int8))
print("MAE:", mae, "MAX:", maxe)

실제 태스크 지표로 검증

  • 분류면 accuracy/F1
  • 검색이면 recall@k
  • 생성이면 BLEU/ROUGE 또는 human eval

로짓 오차가 작아도 softmax 이후 top-1이 바뀌면 품질 이슈가 됩니다. 반드시 서비스 지표로 확인하세요.


6) 성능 측정: “빨라졌는지”를 제대로 재기

양자화 후 속도 개선이 없거나 오히려 느려지는 경우가 있습니다. 이유는 대체로 아래 중 하나입니다.

  • 런타임이 INT8 커널을 못 타는 연산자가 많음
  • 스레드/세션 옵션 최적화 미적용
  • 배치 크기/입력 크기에서 메모리 병목이 큼
  • 모델이 원래부터 작아서 오버헤드가 지배적

간단 벤치마크 코드

import time
import numpy as np
import onnxruntime as ort

def bench(path, n=200, warmup=50):
    so = ort.SessionOptions()
    # 환경에 따라 조정: intra_op_num_threads, graph_optimization_level 등
    sess = ort.InferenceSession(path, sess_options=so, providers=["CPUExecutionProvider"])

    x = np.random.randn(32, 768).astype(np.float32)

    for _ in range(warmup):
        sess.run(None, {"input": x})

    t0 = time.perf_counter()
    for _ in range(n):
        sess.run(None, {"input": x})
    t1 = time.perf_counter()

    return (t1 - t0) * 1000 / n

fp32_ms = bench("mlp_fp32.onnx")
int8_ms = bench("mlp_int8_dynamic.onnx")
print("FP32 ms:", fp32_ms)
print("INT8 ms:", int8_ms)
print("Speedup:", fp32_ms / int8_ms)

측정 시에는 반드시

  • warmup 포함
  • 동일한 provider/스레드 조건
  • 동일한 입력 크기/배치 를 지키세요.

7) Static 양자화(캘리브레이션)로 더 밀어붙이기

Dynamic으로 효과가 제한적이면 static을 고려합니다. 핵심은 “캘리브레이션 데이터가 실제 트래픽 분포를 대표하는가”입니다. 대표성이 떨어지면 특정 구간에서 오차가 튀고 품질 문제가 생깁니다.

ONNX Runtime의 static 양자화는 대략 아래 흐름입니다.

  1. 캘리브레이션 데이터 로더 준비
  2. quantize_static 수행
  3. accuracy 회귀 테스트

아래 코드는 형태만 보여주는 예시입니다.

from onnxruntime.quantization import quantize_static, CalibrationDataReader, QuantType
import numpy as np

class MyDataReader(CalibrationDataReader):
    def __init__(self, n=100):
        self.data = [
            {"input": np.random.randn(32, 768).astype(np.float32)}
            for _ in range(n)
        ]
        self.it = iter(self.data)

    def get_next(self):
        return next(self.it, None)

quantize_static(
    model_input="mlp_fp32.onnx",
    model_output="mlp_int8_static.onnx",
    calibration_data_reader=MyDataReader(n=200),
    weight_type=QuantType.QInt8,
    activation_type=QuantType.QInt8,
)

실전에서는 np.random이 아니라 실제 샘플(로그 기반, 개인정보 마스킹 포함)을 넣어야 합니다.


8) “4배 경량화”가 안 나오는 흔한 이유

1) 양자화된 노드 비율이 낮다

Conv 위주 모델보다 FC/MatMul 위주 모델에서 dynamic 양자화 효과가 잘 나옵니다. 모델 구조에 따라 양자화가 적용되는 연산자가 제한될 수 있습니다.

2) 외부 데이터(External Data) 포맷

ONNX가 큰 텐서를 외부 파일로 분리하는 포맷을 쓰면, 단일 파일 크기만 보고 착시가 생길 수 있습니다. 배포 산출물 전체 크기를 합산해서 비교하세요.

3) 이미 FP16이나 압축이 적용되어 있다

FP16 모델을 INT8로 바꾸면 이론상 2배 감소가 상한인 경우가 많습니다(16bit float → 8bit int). “원본이 FP32인지”부터 확인하세요.

4) 임베딩 테이블이 지배적이고 양자화가 제한된다

추천/검색 모델의 거대한 embedding은 별도 양자화 전략이 필요할 수 있습니다.


9) 배포 관점 팁: 재현 가능한 빌드와 회귀 방지

양자화는 모델 아티팩트 생성 과정 자체가 제품의 일부가 됩니다. 다음을 권장합니다.

  • export 및 quantize 스크립트를 CI에서 고정 실행
  • 캘리브레이션 샘플 버전 관리(해시, 날짜, 샘플링 규칙)
  • FP32 대비 정확도/지연시간 회귀 테스트를 PR 게이트로 운영
  • 문제가 생기면 원인 추적이 가능하도록 “export ONNX”와 “quantized ONNX”를 둘 다 보관

운영에서 성능 문제가 발생했을 때의 원인 추적 접근은 인프라 글이지만 사고 방식은 동일합니다.


10) GPU에서 진짜 성능을 뽑고 싶다면: TensorRT INT8

GPU에서 INT8의 체감 성능을 크게 얻으려면 보통 TensorRT로 엔진을 빌드합니다. 이때는

  • 레이어별 INT8 지원 여부
  • 캘리브레이션(Entropy, MinMax 등)
  • 입력 shape 최적화(동적 shape 범위)
  • 플러그인 필요 여부

같은 이슈가 얽힙니다. ONNX Runtime INT8로는 모델 크기와 CPU 성능을 먼저 챙기고, GPU 최적화가 목표라면 아래 글의 트러블슈팅 포인트를 같이 보는 것을 추천합니다.


마무리: 가장 안전한 실전 진행 순서

  1. PyTorch eval() 기준으로 ONNX export
  2. ONNX Runtime으로 FP32 정확도 스모크
  3. Dynamic INT8로 1차 경량화(파일 크기, 지연시간, 정확도 측정)
  4. 효과가 부족하면 static 캘리브레이션 또는 TensorRT INT8로 확장
  5. 회귀 테스트 자동화로 배포 안정성 확보

이 순서대로 가면 “모델 크기는 줄었는데 품질이 깨졌다” 또는 “속도는 안 빨라졌다” 같은 문제를 단계적으로 분리해서 해결할 수 있습니다.