Huggingface 모델을 ONNX로 변환

  • 카카오톡 공유하기
  • 네이버 블로그 공유하기
  • 네이버 밴드에 공유하기
  • 페이스북 공유하기
  • 트위터 공유하기
  • 링크 복사하기

 

인공지능과 머신러닝이 급속도로 발전하면서, 인공지능 모델을 제공하는 플랫폼들이 생겨났고, 이중Huggingface는 가장 인기있는 플랫폼으로 자리매김했습니다. 다양한 모델 중에 자연어 처리(NLP) 모델을 중심으로 특정한 언어나 도메인에 사전 학습이 된 모델이 있는데 그러나 이러한 모델들을 실제 프로덕션 환경에 배포하고 활용하는 과정에서 여러 문제에 직면하게 됩니다. 이때 ONNX(Open Neural Network Exchange) 변환이 중요한 해결책이 될 수 있습니다.

필자의 회사에서도 Java를 기반으로 하는 소프트웨어가 다수이다보니, 모델 활용에 있어서 인공지능 영역은 python의 결과물을 활용하기엔 제약이 따르기 마련입니다. 이에 이 문제의 우회를 위해 Hugginggface에 다수 올려진 모델 중에서 ONNX로 변환하여 사용하는 방법에 대해서 포스팅하려고 합니다.

ONNX란 무엇인가?

ONNX는 다양한 딥러닝 프레임워크 간의 모델 교환을 가능하게 하는 오픈 포맷입니다. Microsoft와 Facebook이 주도하여 개발한 이 포맷은 PyTorch, TensorFlow, MXNet 등 다양한 프레임워크에서 학습된 모델을 하나의 표준화된 형식으로 변환하여 여러 플랫폼에서 실행할 수 있게 해줍니다.

ONNX로 변환해야 하는 주요 이유

1. 추론 속도 향상

ONNX 런타임은 모델 추론을 위해 최적화된 엔진을 제공합니다. Huggingface의 BERT, GPT, RoBERTa와 같은 대형 모델을 ONNX로 변환하면 다음과 같은 이점이 있습니다:

  • 그래프 최적화: 불필요한 연산을 제거하고 연산을 융합하여 실행 속도 개선
  • 하드웨어 가속: CPU, GPU, TPU 등 다양한 하드웨어에 맞춘 최적화 지원
  • 메모리 사용량 감소: 더 효율적인 메모리 관리로 대형 모델 실행 가능

실제로 BERT와 같은 모델은 ONNX 변환 후 2-5배 빠른 추론 속도를 보여주는 경우가 많습니다.

2. 프레임워크 독립성 확보

Huggingface 모델은 주로 PyTorch나 TensorFlow를 기반으로 합니다. 그러나 프로덕션 환경은 다양한 프레임워크와 플랫폼을 사용할 수 있습니다. ONNX로 변환하면:

  • 프레임워크 종속성 제거: PyTorch나 TensorFlow 설치 없이 모델 실행 가능
  • 다양한 언어 지원: Python 외에도 C++, C#, Java 등에서 모델 사용 가능
  • 버전 호환성 문제 해결: 프레임워크 버전 변경에 따른 문제 감소

3. 경량화 및 배포 용이성

모바일 기기나 엣지 디바이스와 같은 제한된 환경에서 Huggingface 모델을 실행하려면 경량화가 필수적입니다:

  • 모델 크기 축소: 양자화(Quantization)와 같은 기법을 통해 모델 크기 감소
  • 배포 간소화: 단일 파일로 모델 배포 가능
  • 다양한 디바이스 지원: 모바일, 웹 브라우저, IoT 장치 등에 배포 용이

4. 모델 최적화 및 커스터마이징

ONNX 형식으로 변환된 모델은 다양한 최적화 도구를 통해 추가적인 성능 향상이 가능합니다:

  • ONNX Runtime을 통한 자동 최적화
  • 특정 하드웨어에 맞춘 커스텀 최적화 적용
  • 모델 프루닝, 양자화 등 고급 최적화 기법 적용 용이

Huggingface 모델을 ONNX로 변환하는 방법

Huggingface 모델을 ONNX로 변환하는 과정은 비교적 간단합니다. 다음은 BERT 모델을 ONNX로 변환하는 기본 코드 예시입니다:

# 1. 라이브러리 인스톨
# python 3.10.13 기준에서 만들었습니다.
# !pip install optimum[onnxruntime] transformers sentence-transformers torch 
# or please, install dependency library first , 'pip install -r requirements.txt'

from pathlib import Path
from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForFeatureExtraction

# For sentence-transformers pooling:
from sentence_transformers import SentenceTransformer
import torch
import onnxruntime as ort # To verify with ONNX Runtime directly
import numpy as np


print("라이브러리 임포트 완료.")

# === Configuration ===
# 사용할 사전 훈련된 임베딩 모델 이름 (Hugging Face 모델 경로를 그대로 입력해주세요.)
model_name = 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2'

# ONNX 모델을 저장할 경로
onnx_output_dir = Path(f"./{model_name.split('/')[-1]}_onnx")
onnx_output_dir.mkdir(parents=True, exist_ok=True)
onnx_model_path = onnx_output_dir / "model.onnx" # Standard path often used by Optimum

print(f"선택된 모델: {model_name}")
print(f"ONNX 모델 저장 경로: {onnx_output_dir}")

# === Export the model to ONNX using Optimum ===
print("\nOptimum을 사용하여 모델을 ONNX로 변환 시작...")
try:
    # 1. 원본 모델과 토크나이저 로딩
    # Optimum's ORTModel...from_pretrained handles export if export=True
    # It exports the base transformer model. Pooling needs to be handled separately.
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    ort_model = ORTModelForFeatureExtraction.from_pretrained(model_name, export=True)

    # 2. Save the ONNX model and tokenizer configuration
    ort_model.save_pretrained(onnx_output_dir)
    tokenizer.save_pretrained(onnx_output_dir)

    print(f"ONNX 모델 및 토크나이저 설정이 '{onnx_output_dir}'에 성공적으로 저장되었습니다.")
    print(f"저장된 파일 목록: {list(onnx_output_dir.glob('*'))}")

except Exception as e:
    print(f"ONNX 변환 중 오류 발생: {e}")
    print("모델 이름, 라이브러리 설치 상태, 인터넷 연결을 확인하세요.")
    exit()

# === Verification (Optional but Recommended) ===
print("\n--- ONNX 모델 검증 시작 ---")

# Sample text
texts_to_embed = ["이것은 테스트 문장입니다.", "This is a test sentence."]

try:
    # 1. 로딩 ONNX model and tokenizer
    print("저장된 ONNX 모델과 토크나이저 로딩 중...")
    onnx_tokenizer = AutoTokenizer.from_pretrained(onnx_output_dir)
    # Use the same class to load the saved ONNX model
    onnx_ort_model = ORTModelForFeatureExtraction.from_pretrained(onnx_output_dir)
    # Or load directly with onnxruntime session for manual handling
    # ort_session = ort.InferenceSession(str(onnx_model_path))

    print("ONNX 모델 로딩 완료.")

    # 2. 토크나이징 
    inputs = onnx_tokenizer(texts_to_embed, padding=True, truncation=True, return_tensors="np") # Use numpy for ONNX Runtime
    # print("토큰화된 입력:", inputs)

    # 3. ONNX 모델로 추론 실행
    print("ONNX 모델로 추론 실행 중...")
    # Optimum model handles the session internally
    onnx_outputs = onnx_ort_model(**inputs)

    # The output from ORTModelForFeatureExtraction is typically the last hidden state
    last_hidden_state = onnx_outputs.last_hidden_state # shape (batch_size, sequence_length, hidden_size)
    print(f"ONNX 모델 출력 (last_hidden_state) 형태: {last_hidden_state.shape}")

    # 4. Pooling 적용 (IMPORTANT for Sentence Embeddings)
    # Sentence Transformers usually applies mean pooling on top of the transformer output.
    # We need to replicate this manually using the attention mask.
    print("Mean Pooling 적용 중...")
    attention_mask = inputs['attention_mask']
    mask_expanded = np.expand_dims(attention_mask, axis=-1) # Expand mask to match hidden state shape for broadcasting
    sum_embeddings = np.sum(last_hidden_state * mask_expanded, axis=1)
    sum_mask = np.sum(mask_expanded, axis=1)
    sum_mask = np.maximum(sum_mask, 1e-9) # Avoid division by zero
    mean_pooled_embeddings_onnx = sum_embeddings / sum_mask

    print(f"ONNX 모델 + Mean Pooling 최종 임베딩 형태: {mean_pooled_embeddings_onnx.shape}")
    # print("첫 번째 문장 최종 임베딩 (ONNX, 일부):", mean_pooled_embeddings_onnx[0][:5])

    # 5. (Optional) Compare with original Sentence Transformer model
    print("\n(선택) 원본 Sentence Transformer 모델과 결과 비교 중...")
    original_model = SentenceTransformer(model_name)
    original_embeddings = original_model.encode(texts_to_embed, convert_to_numpy=True)

    # print("첫 번째 문장 최종 임베딩 (원본, 일부):", original_embeddings[0][:5])

    # 변환결과를 원본 모델과 비교 검증
    if np.allclose(mean_pooled_embeddings_onnx, original_embeddings, atol=1e-4): # Adjust tolerance if needed
         print(">>> 검증 성공: ONNX 모델 결과가 원본 모델 결과와 매우 유사합니다.")
    else:
         print(">>> 검증 실패: ONNX 모델 결과가 원본 모델 결과와 차이가 있습니다.")
         # Further check: np.max(np.abs(mean_pooled_embeddings_onnx - original_embeddings))

except Exception as e:
    print(f"ONNX 모델 검증 중 오류 발생: {e}")

print("\n--- 스크립트 완료 ---")

requirements.txt 파일입니다.

certifi==2025.4.26
charset-normalizer==3.4.1
coloredlogs==15.0.1
filelock==3.18.0
flatbuffers==25.2.10
fsspec==2025.3.2
git-lfs==1.6
huggingface-hub==0.30.2
humanfriendly==10.0
idna==3.10
Jinja2==3.1.6
joblib==1.4.2
MarkupSafe==3.0.2
mpmath==1.3.0
networkx==3.4.2
numpy==1.26.4
onnx==1.17.0
onnxruntime==1.21.1
optimum==1.24.0
packaging==25.0
pillow==11.2.1
protobuf==6.30.2
PyYAML==6.0.2
regex==2024.11.6
requests==2.32.3
safetensors==0.5.3
scikit-learn==1.6.1
scipy==1.15.2
sentence-transformers==4.1.0
six==1.17.0
spm==0.9.1
sympy==1.14.0
threadpoolctl==3.6.0
tokenizers==0.21.1
torch==2.2.2
tqdm==4.67.1
transformers==4.51.3
typing_extensions==4.13.2
urllib3==2.4.0

ONNX 변환 시 고려사항

Huggingface 모델을 ONNX로 변환할 때 몇 가지 주의해야 할 점이 있습니다:

  • 동적 입력 처리: 가변 길이 시퀀스를 처리하는 방법 고려
  • 모델 아키텍처 호환성: 일부 복잡한 구조는 변환에 제한이 있을 수 있음
  • ONNX 버전 및 Opset 선택: 사용 환경에 맞는 적절한 버전 선택
  • 양자화 전략: INT8, FP16 등 어떤 정밀도로 양자화할지 결정

결론

Huggingface의 다양한 모델을 ONNX로 변환하는 것은 단순한 포맷 변경 이상의 의미를 가집니다. 추론 속도 향상, 프레임워크 독립성, 경량화, 최적화 가능성 등 다양한 이점을 제공하여 실제 프로덕션 환경에서 AI 모델을 효율적으로 활용할 수 있게 해줍니다.

특히 제한된 리소스에서 대형 언어 모델을 운영해야 하는 경우나, 다양한 플랫폼에 모델을 배포해야 하는 경우, ONNX 변환은 거의 필수적인 단계라고 할 수 있습니다. Huggingface와 ONNX의 결합은 최신 AI 모델의 연구 성과를 실제 비즈니스 가치로 전환하는 중요한 다리 역할을 하고 있습니다.


게시됨

카테고리

,

작성자

댓글

답글 남기기

이메일 주소는 공개되지 않습니다. 필수 필드는 *로 표시됩니다