AI 모델 경량화 및 추론 가속 ONNX & TensorRT

import torch
from transformers import CLIPModel

def export_clip_to_onnx(model_id="openai/clip-vit-base-patch32", output_path="clip_image_encoder.onnx"):
    """
    PyTorch로 로드된 CLIP 모델의 Image Encoder 부분만 
    독립적인 ONNX 파일로 추출하여 저장합니다.
    """
    print("1. PyTorch 모델 로드 중...")
    model = CLIPModel.from_pretrained(model_id)
    model.eval() # 반드시 평가 모드로 설정
    
    # 모델에 통과시킬 가짜 입력 데이터 (Dummy Input)
    # Batch Size=1, Color=3(RGB), Width=224, Height=224
    dummy_pixel_values = torch.randn(1, 3, 224, 224)
    
    print("2. ONNX Graph Tracing 시작...")
    # 이미지 인코더만 분리하여 ONNX로 내보내기
    torch.onnx.export(
        model.vision_model,                # 변환할 PyTorch 모델 (vision_model만 추출)
        dummy_pixel_values,                # 모델 구조를 추적하기 위한 더미 입력
        output_path,                       # 저장할 파일명
        export_params=True,                # 학습된 가중치(Weights) 포함 여부
        opset_version=14,                  # 호환성을 위한 ONNX 연산자 버전
        do_constant_folding=True,          # 상수 폴딩 최적화 적용
        input_names=['pixel_values'],      # 입력 노드 이름 정의
        output_names=['image_embeds'],     # 출력 노드 이름 정의
        dynamic_axes={                     # 배치 크기를 고정하지 않고 동적으로 받기 위한 설정
            'pixel_values': {0: 'batch_size'},
            'image_embeds': {0: 'batch_size'}
        }
    )
    print(f"✅ ONNX 변환 완료: {output_path}")

# export_clip_to_onnx()

import onnxruntime as ort
import numpy as np
from PIL import Image
from torchvision import transforms

class ONNXImageEncoder:
    """
    PyTorch 없이 순수 C++ 기반 ONNX Runtime으로 
    이미지 벡터 임베딩을 추출하는 프로덕션 레벨 추론 엔진입니다.
    """
    def __init__(self, onnx_model_path="clip_image_encoder.onnx"):
        # 1. Execution Provider 설정 (하드웨어 가속기 선택)
        # GPU(TensorRT/CUDA)가 있으면 우선 사용, 없으면 CPU 사용
        providers = [
            'TensorrtExecutionProvider', 
            'CUDAExecutionProvider', 
            'CPUExecutionProvider'
        ]
        
        # 2. 인퍼런스 세션 생성 (모델 로드 및 그래프 최적화 수행)
        sess_options = ort.SessionOptions()
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        
        self.session = ort.InferenceSession(
            onnx_model_path, 
            sess_options=sess_options, 
            providers=providers
        )
        
        # 입력 노드와 출력 노드 이름 캐싱
        self.input_name = self.session.get_inputs()[0].name
        
        # numpy 전용 가벼운 전처리 파이프라인 (PIL Image -> Numpy Array)
        self.preprocess = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], 
                                 std=[0.26862954, 0.26130258, 0.27577711])
        ])

    def encode_image(self, pil_image: Image.Image) -> np.ndarray:
        """단일 이미지를 받아 512차원 벡터를 반환합니다."""
        # 1. 이미지 전처리 및 차원 추가 (Batch size=1)
        # PyTorch 텐서 대신 순수 numpy 배열(float32)을 사용합니다.
        input_tensor = self.preprocess(pil_image).unsqueeze(0).numpy().astype(np.float32)
        
        # 2. ONNX Runtime 추론 실행
        # 입력 데이터를 딕셔너리 형태로 전달
        inputs = {self.input_name: input_tensor}
        outputs = self.session.run(None, inputs)
        
        # 3. 임베딩 결과 추출 (L2 정규화 적용)
        embeds = outputs[0]
        embeds = embeds / np.linalg.norm(embeds, axis=1, keepdims=True)
        return embeds

# --- 실무 API 라우터 사용 예시 ---
# onnx_encoder = ONNXImageEncoder("clip_image_encoder.onnx")
# 
# @app.post("/vectorize/image")
# def get_image_vector(file: UploadFile):
#     # 이미지 파일 로드
#     img = Image.open(file.file).convert("RGB")
#     
#     # PyTorch 통과 시 100ms 걸리던 추론이 
#     # ONNX+TensorRT 통과 시 약 10~20ms 수준으로 단축됨
#     vector = onnx_encoder.encode_image(img)
#     
#     return {"vector": vector.tolist()[0]}