beautifulsoups.dev
Published on

PyTorch 8bit PTQ - ONNX+TensorRT로 2배 가속

Authors

서빙 환경에서 "GPU를 더 사지 않고" 지연시간을 줄이는 가장 현실적인 방법 중 하나가 양자화(Quantization) 입니다. 특히 학습을 다시 하지 않는 PTQ(Post-Training Quantization) 는 적용 장벽이 낮고, TensorRT 최적화와 결합하면 FP16 대비 추가적인 이득을 얻는 경우가 많습니다.

이 글에서는 PyTorch 모델을 8bit PTQ로 준비하고, ONNX로 Export한 뒤, TensorRT 엔진으로 빌드해 추론을 2배 수준으로 가속하는 전형적인 파이프라인을 다룹니다. 또한 실제로 자주 터지는 이슈들(동적 shape, calibration 데이터 품질, op 지원 여부, 정확도 하락)을 어떤 순서로 점검하면 좋은지도 함께 정리합니다.

성능 최적화는 결국 "병목을 찾아 줄이는" 과정입니다. 웹 성능에서 Long Task를 추적하듯이 원인을 분리해 접근하면 성공 확률이 올라갑니다. 비슷한 디버깅 사고방식은 Chrome INP 급락 원인 찾기 - Long Task 추적 글과도 맥이 닿습니다.

목표 아키텍처: PyTorch -> ONNX -> TensorRT INT8

전체 흐름은 아래처럼 잡는 것이 가장 안전합니다.

  1. PyTorch에서 모델을 eval() 모드로 고정하고, 입력 shape 정책(고정 혹은 동적)을 결정
  2. ONNX Export 시 불필요한 동적 축을 줄이고, 지원되는 opset을 선택
  3. TensorRT에서 INT8 빌드
    • 가능하면 Q/DQ(Quantize/Dequantize) 기반 또는
    • TensorRT의 PTQ calibration 기반
  4. 정확도 검증(대표 샘플 + edge 케이스)과 성능 측정(latency, throughput, GPU util)

여기서 가장 중요한 분기점은 INT8을 어디서 “결정”하느냐입니다.

  • (A) PyTorch 단계에서 Q/DQ를 삽입해 ONNX에 quantization 정보를 담는 방식
  • (B) ONNX는 FP32 또는 FP16로 내보내고 TensorRT에서 calibration으로 INT8을 찾는 방식

실무에서는 (B)가 더 단순하고, TensorRT 친화적이며, 모델에 따라 성능이 잘 나옵니다. 이 글도 (B) 중심으로 설명합니다.

사전 준비: 버전과 환경을 먼저 고정

양자화는 툴체인 호환성 영향을 크게 받습니다. 다음 조합이 무난합니다.

  • CUDA, cuDNN, TensorRT 버전 호환 확인
  • PyTorch 버전과 ONNX exporter 호환 확인
  • onnx, onnxruntime 버전 고정

Docker로 고정하는 편이 재현성이 좋습니다. 예시는 TensorRT가 포함된 NVIDIA 이미지를 기반으로 합니다.

# 예시: TensorRT가 포함된 NGC 이미지 사용(버전은 환경에 맞게 선택)
docker run --gpus all -it --rm \
  -v "$PWD":/workspace \
  nvcr.io/nvidia/tensorrt:24.01-py3 bash

pip install -U onnx onnxruntime-gpu numpy torch torchvision

1) PyTorch 모델 준비: Export 친화적으로 만들기

ONNX export에서 흔히 실패하는 지점은 다음입니다.

  • torch.no_grad() 누락, eval() 누락
  • 입력이 dict 또는 가변 구조
  • control flow가 많은 모델
  • 동적 shape를 과도하게 열어둠

아래는 최소한의 Export 템플릿입니다.

import torch
import torchvision

def build_model():
    model = torchvision.models.resnet50(weights=None)
    model.eval()
    return model

model = build_model().cuda()

dummy = torch.randn(1, 3, 224, 224, device="cuda")

with torch.no_grad():
    y = model(dummy)
    print(y.shape)

동적 shape를 정말로 써야 할까?

INT8에서 동적 shape는 가능하지만, 프로파일(profile) 설정과 calibration 데이터 준비가 복잡해집니다. 가능하면 아래 전략을 추천합니다.

  • 입력 크기가 고정이면 고정 shape로
  • 배치만 변하면 batch 축만 동적으로
  • 이미지 모델에서 H/W 동적은 정말 필요할 때만

2) ONNX Export: 불필요한 동적 축 최소화

다음 예시는 batch 축만 동적으로 열어두는 Export입니다.

import torch

onnx_path = "resnet50.onnx"

dummy = torch.randn(1, 3, 224, 224, device="cuda")

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

print("saved:", onnx_path)

opset은 무조건 최신이 답이 아니다

  • TensorRT가 특정 opset에서 지원이 더 안정적인 경우가 있습니다.
  • Export가 깨지면 opset을 한 단계 낮추거나, 문제 op를 분리해 확인하세요.

ONNX 모델 검증은 최소한으로라도 해두는 것이 좋습니다.

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

3) TensorRT INT8 PTQ: Calibration 기반으로 엔진 빌드

TensorRT INT8 PTQ의 핵심은 calibration 데이터입니다.

  • 실제 서빙 입력 분포를 대표해야 함
  • 너무 적으면 스케일이 불안정해 정확도 하락
  • 너무 편향되면 특정 클래스/구간에서 오차가 커짐

trtexec로 빠르게 검증하기

가장 빠른 길은 trtexec로 엔진을 만들어 성능과 정확도(대략)를 먼저 보는 것입니다.

고정 shape 예시:

/usr/src/tensorrt/bin/trtexec \
  --onnx=resnet50.onnx \
  --saveEngine=resnet50_int8.engine \
  --int8 \
  --fp16 \
  --workspace=4096 \
  --shapes=input:1x3x224x224 \
  --warmUp=200 \
  --iterations=1000 \
  --useCudaGraph
  • --fp16 는 INT8 빌드에서도 일부 레이어가 FP16로 fallback 될 때 도움이 됩니다.
  • --workspace 는 빌드 탐색 여유를 주며, 너무 작으면 최적화가 제한됩니다.

동적 batch를 쓰는 경우(profile 설정)

동적 batch를 열어뒀다면 아래처럼 min/opt/max 프로파일이 필요합니다.

/usr/src/tensorrt/bin/trtexec \
  --onnx=resnet50.onnx \
  --saveEngine=resnet50_int8.engine \
  --int8 --fp16 \
  --workspace=4096 \
  --minShapes=input:1x3x224x224 \
  --optShapes=input:16x3x224x224 \
  --maxShapes=input:64x3x224x224

Calibration 캐시를 사용해 빌드 반복 비용 줄이기

INT8 calibration은 시간이 꽤 걸릴 수 있습니다. 한 번 만든 calibration 결과를 캐시로 저장해 재사용하세요.

trtexec는 --calib 옵션으로 캐시를 다룹니다(버전에 따라 옵션명이 다를 수 있음). 예시는 다음 형태로 운영합니다.

/usr/src/tensorrt/bin/trtexec \
  --onnx=resnet50.onnx \
  --int8 \
  --calib=calib.cache \
  --saveEngine=resnet50_int8.engine

옵션명이 맞지 않으면 trtexec --help 로 현재 버전의 플래그를 확인하세요.

4) Python에서 TensorRT 엔진 로드 후 추론

실서비스는 trtexec가 아니라 애플리케이션에서 엔진을 로드해 실행합니다. 아래는 TensorRT Python API로 엔진을 로드하고 실행하는 최소 예시입니다.

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

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def load_engine(path: str) -> trt.ICudaEngine:
    with open(path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

engine = load_engine("resnet50_int8.engine")
context = engine.create_execution_context()

# 입력 shape 설정(동적이면 필수)
input_name = engine.get_tensor_name(0)
output_name = engine.get_tensor_name(1)

batch = 1
context.set_input_shape(input_name, (batch, 3, 224, 224))

# 호스트/디바이스 버퍼 준비
inp = np.random.randn(batch, 3, 224, 224).astype(np.float32)

inp_d = cuda.mem_alloc(inp.nbytes)

out_shape = context.get_tensor_shape(output_name)
out = np.empty(out_shape, dtype=np.float32)
out_d = cuda.mem_alloc(out.nbytes)

stream = cuda.Stream()

# 텐서 주소 바인딩
context.set_tensor_address(input_name, int(inp_d))
context.set_tensor_address(output_name, int(out_d))

# 실행
cuda.memcpy_htod_async(inp_d, inp, stream)
context.execute_async_v3(stream_handle=stream.handle)
cuda.memcpy_dtoh_async(out, out_d, stream)
stream.synchronize()

print(out.shape)

주의할 점:

  • 엔진이 INT8이어도 입출력 dtype은 모델/엔진 설정에 따라 FP16 또는 FP32일 수 있습니다.
  • 동적 shape에서는 set_input_shape 호출 후 출력 shape를 다시 조회해야 합니다.

5) "2배 가속"을 현실로 만드는 체크리스트

INT8을 켰는데도 2배가 안 나오는 경우가 흔합니다. 아래 순서로 점검하면 원인 분리가 빠릅니다.

1) 실제로 INT8 커널이 쓰이고 있나

  • 레이어가 많이 FP16 또는 FP32로 fallback 되면 이득이 제한됩니다.
  • trtexec 로그에서 INT8 사용 여부, tactic 선택 결과를 확인하세요.

2) 병목이 GPU가 아니라 전처리/후처리일 수 있다

  • 이미지 decode, resize, normalize가 CPU에서 오래 걸리면 모델만 빨라져도 전체 latency는 그대로입니다.
  • 전처리를 GPU로 옮기거나, 배치/파이프라이닝을 적용하세요.

3) 작은 배치에서는 커널 런치 오버헤드가 지배적

  • batch 1 에서는 FP16과 INT8 차이가 작을 수 있습니다.
  • 목표가 단일 요청 latency인지, throughput인지 먼저 정의하세요.

4) calibration 데이터 품질이 정확도를 좌우

  • 특히 activation 분포가 다양한 모델(검출, 세그, NLP)은 calibration 샘플 수가 중요합니다.
  • 최소 수백에서 수천 샘플을 권장하며, 입력 분포를 대표해야 합니다.

5) 정확도 하락을 줄이는 실전 팁

  • 민감한 레이어는 FP16으로 남기기(혼합 정밀)
  • 입력 정규화/스케일을 학습과 동일하게 맞추기
  • outlier가 많은 경우, calibration 데이터에 outlier를 포함시키기

6) 자주 만나는 문제와 해결 방향

ONNX export는 되는데 TensorRT 빌드가 실패

  • 특정 op가 TensorRT에서 미지원일 수 있습니다.
  • 해결 방향
    • ONNX Graph를 단순화(불필요한 노드 제거)
    • opset 변경
    • TensorRT 플러그인 사용 또는 해당 구간을 분리

결과가 미묘하게 달라서 QA가 통과하지 못함

  • INT8은 근사 연산이라 완전 동일 출력은 기대하기 어렵습니다.
  • 해결 방향
    • 허용 오차 기준을 재정의(예: top-1 일치율, mAP 변화량)
    • 민감 레이어 FP16 고정
    • calibration 데이터 확장

배포 파이프라인에서 캐시가 안 먹어 빌드가 매번 오래 걸림

7) 권장 운영 전략: "측정 가능한" 최적화로 만들기

마지막으로, 2배 가속을 목표로 한다면 다음 지표를 고정해 반복 측정하세요.

  • 모델 단독 latency: P50, P95, P99
  • end-to-end latency: 전처리 + 추론 + 후처리
  • throughput: QPS 또는 images/sec
  • GPU util, SM util, memory bandwidth
  • 정확도 지표: task별 핵심 metric

그리고 변경은 한 번에 하나씩만 적용하세요.

  • FP32 -> FP16
  • FP16 -> INT8(PTQ)
  • 고정 shape -> 동적 shape
  • 전처리 CPU -> GPU

이런 식으로 원인을 분리하면, "왜 빨라졌는지"와 "왜 정확도가 떨어졌는지"를 설명 가능한 상태로 유지할 수 있습니다.

마무리

PyTorch 모델을 8bit PTQ로 양자화하고 ONNX를 거쳐 TensorRT로 실행하면, 모델/입력/배치 조건이 맞는 경우 FP16 대비 체감 2배 수준의 가속도 충분히 가능합니다. 다만 성공의 핵심은 --int8 플래그가 아니라,

  • Export 가능한 그래프 형태
  • 대표성 있는 calibration 데이터
  • 동적 shape와 profile의 절제
  • end-to-end 병목 분리

에 달려 있습니다.

원하시면 사용 중인 모델 종류(분류/검출/세그/NLP), 입력 shape 정책(고정 또는 동적), 목표 지표(P99 latency 또는 throughput)에 맞춰 calibration 샘플 구성과 TensorRT 빌드 플래그를 더 구체적으로 튜닝하는 체크리스트도 이어서 정리해드릴 수 있습니다.

Discuss on Twitter • View on GitHub