ONNX Runtime과 TensorRT로 컴퓨터비전 모델 경량화하기: Edge AI 배포 실전 가이드

import torch
import torchvision.models as models
import onnxruntime as ort
import numpy as np

# 1. PyTorch 모델 로드
model = models.resnet50(pretrained=True)
model.eval()

# 2. ONNX로 변환
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "resnet50.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

# 3. ONNX Runtime 세션 생성 (CPU)
sess = ort.InferenceSession("resnet50.onnx", providers=["CPUExecutionProvider"])

# 4. 추론 실행
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
output = sess.run(None, {"input": input_data})
print(f"출력 shape: {output[0].shape}")  # (1, 1000)

# CUDA Execution Provider로 GPU 가속
sess_options = ort.SessionOptions()
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

sess_gpu = ort.InferenceSession(
    "resnet50.onnx",
    sess_options,
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
)

# FP16 변환 (NVIDIA GPU)
import onnx
from onnxconverter_common import float16

model_fp32 = onnx.load("resnet50.onnx")
model_fp16 = float16.convert_float_to_float16(model_fp32)
onnx.save(model_fp16, "resnet50_fp16.onnx")

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

# Calibration 데이터셋 준비
class ImageCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, data_loader, cache_file):
        super().__init__()
        self.data_loader = data_loader
        self.cache_file = cache_file
        self.batch_size = 8
        self.current_index = 0
        # GPU 메모리 할당
        self.device_input = cuda.mem_alloc(self.batch_size * 3 * 224 * 224 * 4)

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        if self.current_index < len(self.data_loader):
            batch = next(iter(self.data_loader)).numpy()
            cuda.memcpy_htod(self.device_input, batch)
            self.current_index += 1
            return [int(self.device_input)]
        return None

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

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

# TensorRT 엔진 빌드 (INT8)
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

with open("resnet50.onnx", "rb") as f:
    parser.parse(f.read())

config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = ImageCalibrator(calib_loader, "calibration.cache")

engine = builder.build_serialized_network(network, config)
with open("resnet50_int8.trt", "wb") as f:
    f.write(engine)

# 동적 배치 사이즈 설정
profile = builder.create_optimization_profile()
profile.set_shape(
    "input",
    (1, 3, 224, 224),    # min
    (4, 3, 224, 224),    # opt
    (16, 3, 224, 224)    # max
)
config.add_optimization_profile(profile)

# 양자화 레벨별 설정
quantization_levels = {
    "FP32": {"accuracy": 100, "size": 100, "speed": 1.0},
    "FP16": {"accuracy": 99.5, "size": 50, "speed": 2.0},
    "INT8": {"accuracy": 98.5, "size": 25, "speed": 3.5},
    "Mixed (FP16+INT8)": {"accuracy": 99.0, "size": 30, "speed": 3.0}
}

import cv2
import tensorrt as trt
import numpy as np

class TRTInference:
    def __init__(self, engine_path):
        self.logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, "rb") as f:
            self.runtime = trt.Runtime(self.logger)
            self.engine = self.runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()

    def infer(self, image):
        # 전처리
        input_data = cv2.resize(image, (640, 640))
        input_data = input_data.astype(np.float32) / 255.0
        input_data = np.transpose(input_data, (2, 0, 1))[np.newaxis, :]

        # 추론
        self.context.execute_v2([input_data.ctypes.data, output.ctypes.data])
        return output

# 실시간 스트림 처리
model = TRTInference("yolov8n_int8.trt")
cap = cv2.VideoCapture(0)

while True:
    ret, frame = cap.read()
    detections = model.infer(frame)
    # 후처리 및 시각화...
    cv2.imshow("Detection", frame)
    if cv2.waitKey(1) == ord('q'):
        break

# ONNX Runtime 메모리 설정
sess_options.enable_mem_pattern = True
sess_options.enable_cpu_mem_arena = True

# TensorRT Workspace 제한 (Jetson용)
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 28)  # 256MB

sess_options.intra_op_num_threads = 4

pip install onnx-simplifier
python -m onnxsim model.onnx model_simplified.onnx