Skip to main content

​
ONNX / TensorRT 변환

PyTorch 모델을 ONNX(Open Neural Network Exchange) 형식으로 내보내고, NVIDIA TensorRT로 최적화하여 추론 속도를 극대화하는 과정을 실습합니다.
1
ONNX Export (일반 PyTorch 모델)
2
import torch
import timm

# 모델 준비
model = timm.create_model('efficientnet_b0', pretrained=True, num_classes=5)
model.eval()

# 더미 입력
dummy_input = torch.randn(1, 3, 224, 224)

# ONNX Export
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,
)
print("ONNX 내보내기 완료")
3
ONNX 모델 검증
4
import onnx
import onnxruntime as ort
import numpy as np

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

# ONNX Runtime 추론
session = ort.InferenceSession('model.onnx', providers=['CUDAExecutionProvider'])

input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = session.run(None, {'input': input_data})
print(f"출력 형태: {outputs[0].shape}")
5
YOLO 모델 ONNX/TensorRT Export
6
from ultralytics import YOLO

model = YOLO('best.pt')

# ONNX Export
model.export(format='onnx', imgsz=640, simplify=True)

# TensorRT Export (NVIDIA GPU 필요)
model.export(format='engine', imgsz=640, half=True)  # FP16
7
TensorRT 직접 변환
8
# trtexec CLI 사용 (NVIDIA TensorRT 설치 필요)
# FP16 최적화
# trtexec --onnx=model.onnx --saveEngine=model.engine --fp16

# INT8 최적화 (캘리브레이션 필요)
# trtexec --onnx=model.onnx --saveEngine=model_int8.engine --int8
#         --calib=calibration_data/
9
속도 벤치마크
10
import time

def benchmark_latency(session, input_data, n_runs=100):
    """추론 지연시간을 벤치마크합니다."""
    # 워밍업
    for _ in range(10):
        session.run(None, {'input': input_data})

    # 측정
    start = time.perf_counter()
    for _ in range(n_runs):
        session.run(None, {'input': input_data})
    elapsed = (time.perf_counter() - start) / n_runs * 1000

    print(f"평균 지연시간: {elapsed:.2f} ms ({1000/elapsed:.1f} FPS)")

input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
benchmark_latency(session, input_data)

​
최적화 파이프라인

​
런타임 비교

런타임대상 하드웨어속도설치 난이도
PyTorchGPU/CPU기준낮음
ONNX RuntimeGPU/CPU1.5~2x낮음
TensorRTNVIDIA GPU3~5x중간
OpenVINOIntel CPU2~3x중간
opset_version을 높이거나(17 이상), 커스텀 연산을 ONNX 호환 연산으로 교체하세요. torch.onnx.export에서 verbose=True로 변환 과정을 확인할 수 있습니다.
NVIDIA GPU가 있으면 TensorRT가 가장 빠릅니다. 다양한 하드웨어를 지원해야 하거나 설치가 간편해야 하면 ONNX Runtime을 선택하세요.