파이썬 ONNX Runtime로 INT8 양자화 실전 가이드

pip install onnx onnxruntime onnxruntime-tools numpy

import time
import numpy as np
import onnxruntime as ort


def make_session(model_path: str, threads: int = 1):
    so = ort.SessionOptions()
    so.intra_op_num_threads = threads
    so.inter_op_num_threads = 1
    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

    providers = ["CPUExecutionProvider"]
    return ort.InferenceSession(model_path, sess_options=so, providers=providers)


def benchmark(session: ort.InferenceSession, feed: dict, warmup: int = 10, iters: int = 100):
    for _ in range(warmup):
        session.run(None, feed)

    t0 = time.perf_counter()
    for _ in range(iters):
        session.run(None, feed)
    t1 = time.perf_counter()

    return (t1 - t0) * 1000.0 / iters


# 예시: 입력 하나짜리 모델이라고 가정
fp32_model = "model_fp32.onnx"

sess = make_session(fp32_model, threads=4)
input_name = sess.get_inputs()[0].name
input_shape = sess.get_inputs()[0].shape

# 동적 배치/시퀀스 길이면 None이 있을 수 있으니 적절히 채워야 함
shape = [d if isinstance(d, int) else 1 for d in input_shape]
feed = {input_name: np.random.randn(*shape).astype(np.float32)}

lat_ms = benchmark(sess, feed)
print(f"FP32 latency: {lat_ms:.3f} ms")

from onnxruntime.quantization import quantize_dynamic, QuantType

fp32_model = "model_fp32.onnx"
int8_dynamic_model = "model_int8_dynamic.onnx"

quantize_dynamic(
    model_input=fp32_model,
    model_output=int8_dynamic_model,
    weight_type=QuantType.QInt8,  # 보통 QInt8 또는 QUInt8 중 선택
)

print("saved:", int8_dynamic_model)

import numpy as np
from onnxruntime.quantization import CalibrationDataReader


class NumpyCalibrationDataReader(CalibrationDataReader):
    def __init__(self, input_name: str, samples: list[np.ndarray]):
        self.input_name = input_name
        self.samples = samples
        self._iter = iter(self.samples)

    def get_next(self):
        try:
            x = next(self._iter)
        except StopIteration:
            return None
        return {self.input_name: x}

    def rewind(self):
        self._iter = iter(self.samples)

from onnxruntime.quantization import (
    quantize_static,
    QuantFormat,
    QuantType,
    CalibrationMethod,
)
import onnxruntime as ort

fp32_model = "model_fp32.onnx"
int8_static_model = "model_int8_static.onnx"

# 세션으로 입력 이름을 얻어온 뒤, 캘리브레이션 샘플을 준비
sess = ort.InferenceSession(fp32_model, providers=["CPUExecutionProvider"])
input_name = sess.get_inputs()[0].name
input_shape = sess.get_inputs()[0].shape
shape = [d if isinstance(d, int) else 1 for d in input_shape]

# 실제로는 서비스 트래픽에서 샘플링한 전처리 결과를 쓰는 것이 이상적
samples = [np.random.randn(*shape).astype(np.float32) for _ in range(200)]
calib_reader = NumpyCalibrationDataReader(input_name, samples)

quantize_static(
    model_input=fp32_model,
    model_output=int8_static_model,
    calibration_data_reader=calib_reader,
    quant_format=QuantFormat.QDQ,           # 보통 QDQ 권장
    activation_type=QuantType.QInt8,
    weight_type=QuantType.QInt8,
    calibrate_method=CalibrationMethod.MinMax,  # Entropy 등도 비교 가능
)

print("saved:", int8_static_model)

import numpy as np
import onnxruntime as ort


def run_outputs(model_path: str, feed: dict):
    sess = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
    return sess.run(None, feed)


def cosine_sim(a: np.ndarray, b: np.ndarray, eps: float = 1e-9) -> float:
    a = a.reshape(-1).astype(np.float32)
    b = b.reshape(-1).astype(np.float32)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + eps))


fp32 = "model_fp32.onnx"
int8 = "model_int8_dynamic.onnx"

sess = ort.InferenceSession(fp32, providers=["CPUExecutionProvider"])
input_name = sess.get_inputs()[0].name
shape = [d if isinstance(d, int) else 1 for d in sess.get_inputs()[0].shape]

x = np.random.randn(*shape).astype(np.float32)
feed = {input_name: x}

out_fp32 = run_outputs(fp32, feed)[0]
out_int8 = run_outputs(int8, feed)[0]

print("cosine:", cosine_sim(out_fp32, out_int8))
print("max_abs_diff:", float(np.max(np.abs(out_fp32 - out_int8))))