PyTorch→ONNX→TensorRT INT8 양자화 실전 가이드

서빙 환경에서 GPU 추론 성능을 더 끌어올리고 싶다면, PyTorch → ONNX → TensorRT 파이프라인은 사실상 표준에 가깝습니다. 특히 INT8 양자화는 메모리 대역폭과 연산량을 크게 줄여 지연시간을 낮추는 데 효과적이지만, 정확도 저하와 캘리브레이션 데이터 준비, 동적 shape 처리 같은 실무 이슈가 따라옵니다.

이 글은 "일단 돌아가게" 수준이 아니라, 실제 배포를 염두에 둔 형태로 INT8 빌드와 검증까지 한 번에 정리합니다.

디버깅 팁: 최적화 작업은 실험을 반복하므로 CI에서 캐시가 중요합니다. ONNX/TensorRT 빌드 산출물 캐시 전략은 GitHub Actions 캐시로 Node.js CI 2배 빠르게, 실패 디버깅에서 아이디어를 얻을 수 있습니다.

전체 흐름과 선택지

목표 파이프라인

1. PyTorch에서 추론용 모델 준비(eval, no_grad, 고정된 입력 규약)
2. ONNX export(가능하면 opset 고정, dynamic axes 정의)
3. ONNX 정리(상수 폴딩, shape 추론, 그래프 단순화)
4. TensorRT 엔진 빌드
   • FP16 먼저 성공시키기
   • INT8은 캘리브레이션 포함해 빌드
5. 정확도/성능 검증

INT8 양자화 방식: PTQ vs QAT

• PTQ(Post-Training Quantization): 학습 없이 캘리브레이션 데이터만으로 스케일을 추정. 가장 흔한 실전 루트.
• QAT(Quantization-Aware Training): 학습 때부터 양자화를 고려. 정확도는 더 잘 나오지만 비용이 큼.

이 글은 PTQ 기반 INT8을 중심으로 설명합니다.

사전 준비: 버전 호환성과 환경

TensorRT는 CUDA, cuDNN, GPU 아키텍처, ONNX opset에 민감합니다. 실무에서는 아래 원칙이 안전합니다.

• torch, onnx, onnxruntime-gpu, tensorrt 버전을 한 세트로 고정
• opset은 너무 최신을 피하고 TensorRT가 잘 지원하는 범위를 선택
• 먼저 FP16 엔진이 정상 동작하는지 확인한 뒤 INT8로 확장

권장 진단 루틴은 trtexec로 ONNX를 바로 빌드해 보고, 실패하면 ONNX 그래프를 단순화하거나 export 옵션을 조정하는 방식입니다.

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

핵심은 "TensorRT가 좋아하는" ONNX를 만드는 것입니다.

• eval() 필수(드롭아웃/배치정규화 동작 고정)
• 입력 shape를 명확히
• 가능하면 torch.jit.trace로 경로를 단순화하거나 export 시 do_constant_folding=True

아래 예시는 이미지 분류 모델을 ONNX로 export하는 기본 형태입니다.

import torch

def export_onnx(model, onnx_path="model.onnx", opset=17):
    model.eval()
    dummy = torch.randn(1, 3, 224, 224, device="cpu")

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

# 사용 예
# export_onnx(model, "resnet.onnx", opset=17)

ONNX export에서 자주 터지는 포인트

• aten:: 연산이 남아 TensorRT가 못 읽는 경우
• Resize, GridSample, NonMaxSuppression 등 특정 연산이 버전에 따라 제한
• dynamic shape를 과하게 열어두면 최적화가 약해짐

실무 팁은 "batch만 dynamic"으로 두고, 나머지 H/W는 고정하는 것입니다. 특히 비전 모델은 입력 해상도를 고정할수록 TensorRT 최적화가 잘 됩니다.

2) ONNX 그래프 점검과 단순화

ONNX를 export했다고 끝이 아닙니다. 최소한 다음은 확인하세요.

• ONNX 체크
• shape inference
• 가능하면 simplifier 적용

python -c "import onnx; m=onnx.load('model.onnx'); onnx.checker.check_model(m); print('OK')"

onnx-simplifier를 쓰면 불필요한 노드를 줄이고 TensorRT 호환성이 좋아지는 경우가 많습니다.

python -m pip install onnxsim
python -m onnxsim model.onnx model.sim.onnx

주의: simplifier가 그래프를 바꾸면서 미세한 수치 차이가 생길 수 있으니, 반드시 정확도 검증을 함께 하세요.

3) TensorRT에서 FP16 엔진 먼저 만들기

INT8로 바로 가면 실패 원인을 분리하기 어렵습니다. 먼저 FP16 빌드가 성공하고 출력이 정상인지 확인합니다.

trtexec는 가장 빠른 스모크 테스트 도구입니다.

trtexec \
  --onnx=model.sim.onnx \
  --saveEngine=model.fp16.engine \
  --fp16 \
  --workspace=4096 \
  --verbose

여기서부터는 로그가 중요합니다.

• 특정 레이어가 플러그인으로 떨어지는지
• precision fallback이 발생하는지
• dynamic shape면 optimization profile이 어떻게 잡히는지

4) INT8 PTQ 캘리브레이션: 데이터가 성능과 정확도를 좌우

TensorRT INT8 PTQ는 캘리브레이션 데이터로 activation 분포를 관측해 스케일을 정합니다.

캘리브레이션 데이터 원칙:

• 실서빙 입력 분포를 최대한 반영
• 너무 적으면 스케일이 불안정
• 너무 특이한 샘플을 넣으면 오히려 정확도 악화

실무적으로는 클래스/상황을 골고루 포함한 수백~수천 샘플을 권장합니다.

캘리브레이션 방식 선택

TensorRT에는 대표적으로 아래가 있습니다.

• Entropy(일반적으로 정확도 좋음)
• MinMax(분포가 단순할 때 빠르고 안정적)

Python API로 IInt8EntropyCalibrator2를 구현하는 방식이 흔합니다.

5) TensorRT Python API로 INT8 엔진 빌드

아래 코드는 "동적 배치"를 가진 ONNX를 로드하고, INT8 캘리브레이터를 붙여 엔진을 빌드하는 최소 예시입니다. 전처리와 데이터 로더 부분은 프로젝트에 맞게 바꿔야 합니다.

주의: 본문에서 부등호 문자가 그대로 노출되면 MDX에서 문제가 될 수 있어, 제네릭/화살표 등은 모두 인라인 코드로 표기합니다.

import os
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

TRT_LOGGER = trt.Logger(trt.Logger.INFO)

class ImageBatchCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, batches, input_name="input", cache_file="calib.cache"):
        super().__init__()
        self.batches = batches
        self.input_name = input_name
        self.cache_file = cache_file
        self.batch_idx = 0

        first = self.batches[0]
        self.device_input = cuda.mem_alloc(first.nbytes)

    def get_batch_size(self):
        return int(self.batches[0].shape[0])

    def get_batch(self, names):
        if self.batch_idx >= len(self.batches):
            return None
        batch = self.batches[self.batch_idx]
        self.batch_idx += 1
        cuda.memcpy_htod(self.device_input, batch)
        return [int(self.device_input)]

    def read_calibration_cache(self):
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                return f.read()
        return None

    def write_calibration_cache(self, cache):
        with open(self.cache_file, "wb") as f:
            f.write(cache)


def build_int8_engine(
    onnx_path="model.sim.onnx",
    engine_path="model.int8.engine",
    input_name="input",
    max_batch=16,
    workspace_mb=4096,
    calib_batches=None,
):
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    )
    parser = trt.OnnxParser(network, TRT_LOGGER)

    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_mb * 1024 * 1024)

    with open(onnx_path, "rb") as f:
        if not parser.parse(f.read()):
            for i in range(parser.num_errors):
                print(parser.get_error(i))
            raise RuntimeError("ONNX parse failed")

    profile = builder.create_optimization_profile()
    # 예시: NCHW, batch만 dynamic
    profile.set_shape(input_name, (1, 3, 224, 224), (4, 3, 224, 224), (max_batch, 3, 224, 224))
    config.add_optimization_profile(profile)

    config.set_flag(trt.BuilderFlag.INT8)

    if calib_batches is None:
        raise ValueError("calib_batches is required for INT8 PTQ")

    calibrator = ImageBatchCalibrator(calib_batches, input_name=input_name)
    config.int8_calibrator = calibrator

    serialized = builder.build_serialized_network(network, config)
    if serialized is None:
        raise RuntimeError("Engine build failed")

    with open(engine_path, "wb") as f:
        f.write(serialized)

    return engine_path


# calib_batches 예시 생성: (B, 3, 224, 224) float32, 전처리된 입력
# 실제로는 데이터셋에서 로드해 np.array로 만드세요.
calib_batches = [
    (np.random.rand(4, 3, 224, 224).astype(np.float32) * 2.0 - 1.0)
    for _ in range(50)
]

# build_int8_engine(calib_batches=calib_batches)

캘리브레이션 입력 dtype과 스케일

• TensorRT 캘리브레이션은 보통 float32 입력을 기대합니다.
• 전처리(정규화, mean/std, resize)가 학습/서빙과 동일해야 합니다.
• uint8 원본 이미지를 그대로 넣고 싶다면, 네트워크 입력이 그 스케일을 반영하도록 ONNX 그래프에 전처리를 포함시키거나, 캘리브레이터에서 동일 전처리를 수행해야 합니다.

6) 정확도 검증: "엔진이 만들어졌다"와 "정답을 맞춘다"는 다릅니다

INT8은 FP16 대비 정확도 손실이 발생할 수 있습니다. 실무에서는 아래 3단계를 권장합니다.

1. PyTorch FP32 vs ONNXRuntime FP32 출력 비교
2. ONNXRuntime FP32 vs TensorRT FP16 출력 비교
3. TensorRT FP16 vs TensorRT INT8 출력 비교

출력 비교는 완전 일치가 아니라 허용 오차 기반으로 봅니다.

import numpy as np

def compare(a, b, name="tensor"):
    a = np.asarray(a)
    b = np.asarray(b)
    max_abs = np.max(np.abs(a - b))
    mean_abs = np.mean(np.abs(a - b))
    print(f"{name}: max_abs={max_abs:.6f}, mean_abs={mean_abs:.6f}")

분류 모델이면 최종적으로 top-1/top-5 정확도를 기준으로 비교하는 것이 가장 명확합니다.

7) 성능 측정: 지연시간은 프로파일과 배치에 따라 달라진다

INT8이 항상 이기는 것은 아닙니다.

• 작은 배치에서 커널 런치 오버헤드가 지배적이면 FP16과 차이가 작을 수 있음
• 일부 레이어가 INT8로 내려가지 못하고 FP16/FP32로 fallback되면 이득이 줄어듦
• 메모리 바운드 모델에서 효과가 더 큼

trtexec로 동일 조건에서 FP16과 INT8을 비교하세요.

trtexec --loadEngine=model.fp16.engine --warmUp=200 --duration=10 --useCudaGraph
trtexec --loadEngine=model.int8.engine --warmUp=200 --duration=10 --useCudaGraph

측정 시에는 입력 shape, batch, 스트림 수를 고정하고, 워밍업을 충분히 줘야 합니다.

8) 실전에서 자주 만나는 함정과 해결책

1) INT8 빌드는 되는데 정확도가 크게 떨어진다

점검 순서:

• 캘리브레이션 데이터가 실제 분포를 반영하는지
• 전처리가 학습과 완전히 동일한지
• dynamic shape를 과도하게 열어두지 않았는지
• 특정 레이어가 precision-sensitive인지

대응:

• 캘리브레이션 데이터 확대 및 대표성 개선
• MinMax 방식으로 바꿔보기(분포가 단순한 모델에서 안정적인 경우가 있음)
• 일부 레이어만 FP16 유지하는 "혼합 정밀도" 전략 검토

2) 동적 shape에서 최적화가 약해진다

TensorRT는 optimization profile 범위가 넓을수록 보수적으로 최적화합니다.

• profile을 1개만 크게 잡기보다, 자주 쓰는 배치/해상도 조합별로 profile을 여러 개 두는 것이 유리할 때가 많습니다.

3) ONNX 파서 에러가 난다

• 먼저 FP16 trtexec로 에러 재현
• 문제 노드 주변을 확인
• ONNX simplifier 적용
• export 시 opset 변경, 특정 연산 대체

4) 배포 파이프라인에서 빌드가 느리고 불안정하다

엔진 빌드는 환경 의존성이 강합니다. 운영에서는 다음을 권장합니다.

• 빌드 컨테이너와 런타임 컨테이너를 분리
• GPU 아키텍처가 다른 곳에서 빌드한 엔진을 그대로 가져오지 않기
• CI에서 빌드 산출물 캐시(ONNX, 엔진, 캘리브레이션 캐시)를 적극 활용

CI 캐시 운영 아이디어는 GitHub Actions 캐시로 Node.js CI 2배 빠르게, 실패 디버깅도 참고할 만합니다.

9) 운영 체크리스트

• 입력 전처리 완전 동일성(학습, 캘리브레이션, 서빙)
• FP16 엔진을 기준선으로 확보
• INT8은 캘리브레이션 캐시를 아티팩트로 보관
• 지연시간은 p50, p95, p99로 보고, 배치/동시성 조건을 명시
• 드라이버/CUDA 업그레이드 시 엔진 재빌드 필요 여부 확인

마무리

PyTorch → ONNX → TensorRT로 INT8까지 가는 길은 도구 체인이 길어서 한 번에 매끄럽게 성공하기 어렵습니다. 하지만 순서를 잘 나누면(ONNX export 안정화 → FP16 성공 → INT8 캘리브레이션/검증) 재현 가능하고 운영 가능한 파이프라인을 만들 수 있습니다.

다음 단계로는 QAT 적용, 레이어별 정밀도 제어, 다중 optimization profile 운영, 그리고 실제 트래픽 분포를 반영한 캘리브레이션 데이터 자동 수집까지 확장해보면 좋습니다.