ONNX로 내보내기

프로덕션 환경에 🤗 Transformers 모델을 배포할 때에는 특수 런타임 및 하드웨어 위에 올리고 실행할 수 있도록 직렬화된 형식으로 내보내기를 권장합니다. 이 가이드에서는 🤗 Transformers 모델을 ONNX (Open Neural Network eXchange)로 내보내는 방법을 안내합니다.

ONNX는 딥러닝 모델을 표현하기 위한 공통 파일 형식과 연산자들을 정의하는 개방형 표준으로써 PyTorch, TensorFlow 등 다양한 프레임워크에서 지원됩니다. 모델을 ONNX 형식으로 내보내면, (보통 중간 표현 (Intermediate Representation; IR)이라고 불리는) 계산 그래프가 구성됩니다. 계산 그래프는 신경망을 통해 데이터가 흐르는 방식, 즉 어떤 연산이 어느 부분에 사용되었는지를 나타냅니다.

표준 연산 및 데이터 형식을 사용하여 그래프를 노출하기 때문에 ONNX를 사용하면 프레임워크 간 전환이 쉬워집니다. 예를 들어, PyTorch에서 훈련된 모델은 ONNX 형식으로 내보낸 뒤, TensorFlow에서 가져올 수 있습니다. 물론 그 반대도 가능합니다.

🤗 Transformers는 모델 체크포인트를 ONNX 그래프로 변환할 수 있게 해주는 transformers.onnx 패키지를 제공합니다. 이걸 가능케 하는 구성 객체는 여러 모델 아키텍처를 대상으로 미리 제작되어 있으며, 다른 아키텍처로도 쉽게 확장할 수 있도록 설계되었습니다.

미리 제작된 구성에는 다음 아키텍처가 포함됩니다:

• ALBERT
• BART
• BEiT
• BERT
• BigBird
• BigBird-Pegasus
• Blenderbot
• BlenderbotSmall
• BLOOM
• CamemBERT
• Chinese-CLIP
• CLIP
• CodeGen
• Conditional DETR
• ConvBERT
• ConvNeXT
• Data2VecText
• Data2VecVision
• DeBERTa
• DeBERTa-v2
• DeiT
• DETR
• DistilBERT
• EfficientNet
• ELECTRA
• ERNIE
• FlauBERT
• GPT Neo
• GPT-J
• GPT-Sw3
• GroupViT
• I-BERT
• ImageGPT
• LayoutLM
• LayoutLMv3
• LeViT
• Longformer
• LongT5
• M2M100
• Marian
• mBART
• MEGA
• MobileBERT
• MobileNetV1
• MobileNetV2
• MobileViT
• MT5
• OpenAI GPT-2
• OWL-ViT
• Perceiver
• PLBart
• PoolFormer
• RemBERT
• ResNet
• RoBERTa
• RoBERTa-PreLayerNorm
• RoFormer
• SegFormer
• SqueezeBERT
• Swin Transformer
• T5
• Table Transformer
• Vision Encoder decoder
• ViT
• Whisper
• X-MOD
• XLM
• XLM-RoBERTa
• XLM-RoBERTa-XL
• YOLOS

앞으로의 두 섹션에서는 아래 내용을 살펴보겠습니다:

• transformers.onnx 패키지를 사용하여 지원되는 모델 내보내기
• 지원되지 않는 아키텍처를 위해 사용자 정의 모델 내보내기

모델을 ONNX로 내보내기

🤗 Transformers 모델을 ONNX로 내보내려면 먼저 몇 가지 추가 종속성을 설치해야합니다:

pip install transformers[onnx]

transformers.onnx 패키지는 다음과 같이 Python 모듈로 사용할 수 있습니다:

python -m transformers.onnx --help

usage: Hugging Face Transformers ONNX exporter [-h] -m MODEL [--feature {causal-lm, ...}] [--opset OPSET] [--atol ATOL] output

positional arguments:
  output                Path indicating where to store generated ONNX model.

optional arguments:
  -h, --help            show this help message and exit
  -m MODEL, --model MODEL
                        Model ID on huggingface.co or path on disk to load model from.
  --feature {causal-lm, ...}
                        The type of features to export the model with.
  --opset OPSET         ONNX opset version to export the model with.
  --atol ATOL           Absolute difference tolerance when validating the model.

다음과 같이 미리 제작된 구성을 사용하여 체크포인트를 내보낼 수 있습니다:

python -m transformers.onnx --model=distilbert-base-uncased onnx/

다음과 같은 로그가 표시되어야합니다:

Validating ONNX model...
        -[✓] ONNX model output names match reference model ({'last_hidden_state'})
        - Validating ONNX Model output "last_hidden_state":
                -[✓] (2, 8, 768) matches (2, 8, 768)
                -[✓] all values close (atol: 1e-05)
All good, model saved at: onnx/model.onnx

이렇게 --model 인수로 정의된 체크포인트의 ONNX 그래프를 내보냅니다. 예시에서는 distilbert-base-uncased이지만, Hugging Face Hub에서 가져왔거나 로컬에 저장된 체크포인트들 모두 가능합니다.

결과로 나온 model.onnx 파일은 ONNX 표준을 지원하는 다양한 가속기 중 하나에서 실행할 수 있습니다. 예를 들어, 다음과 같이 ONNX Runtime에서 모델을 가져오고 실행할 수 있습니다:

>>> from transformers import AutoTokenizer
>>> from onnxruntime import InferenceSession

>>> tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
>>> session = InferenceSession("onnx/model.onnx")
>>> # ONNX Runtime expects NumPy arrays as input
>>> inputs = tokenizer("Using DistilBERT with ONNX Runtime!", return_tensors="np")
>>> outputs = session.run(output_names=["last_hidden_state"], input_feed=dict(inputs))

["last_hidden_state"]와 같은 필요한 출력 이름은 각 모델의 ONNX 구성을 살펴보면 얻을 수 있습니다. 예를 들어, DistilBERT의 경우 다음과 같습니다:

>>> from transformers.models.distilbert import DistilBertConfig, DistilBertOnnxConfig

>>> config = DistilBertConfig()
>>> onnx_config = DistilBertOnnxConfig(config)
>>> print(list(onnx_config.outputs.keys()))
["last_hidden_state"]

Hub의 TensorFlow 체크포인트의 경우에도 과정은 동일합니다. 예를 들어, 다음과 같이 Keras organization에서 TensorFlow 체크포인트를 내보낼 수 있습니다:

python -m transformers.onnx --model=keras-io/transformers-qa onnx/

로컬에 저장된 모델을 내보내려면 모델의 가중치 및 토크나이저 파일이 저장된 디렉토리가 필요합니다. 예를 들어, 다음과 같이 체크포인트를 가져오고 저장할 수 있습니다:

>>> from transformers import AutoTokenizer, AutoModelForSequenceClassification

>>> # Load tokenizer and PyTorch weights form the Hub
>>> tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
>>> pt_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")
>>> # Save to disk
>>> tokenizer.save_pretrained("local-pt-checkpoint")
>>> pt_model.save_pretrained("local-pt-checkpoint")

python -m transformers.onnx --model=local-pt-checkpoint onnx/

>>> from transformers import AutoTokenizer, TFAutoModelForSequenceClassification

>>> # Load tokenizer and TensorFlow weights from the Hub
>>> tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
>>> tf_model = TFAutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")
>>> # Save to disk
>>> tokenizer.save_pretrained("local-tf-checkpoint")
>>> tf_model.save_pretrained("local-tf-checkpoint")

python -m transformers.onnx --model=local-tf-checkpoint onnx/

다른 모델 작업에 대한 기능 선택

다른 유형의 태스크에 맞춰서 모델을 내보낼 수 있도록 미리 제작된 구성마다 일련의 기능이 포함되어 있습니다. 아래 표에 나와 있는대로 각 기능은 다른 AutoClass와 연관되어 있습니다.

각 구성에서 FeaturesManager를 통해 지원되는 기능 목록을 찾을 수 있습니다. 예를 들어, DistilBERT의 경우 다음과 같습니다:

>>> from transformers.onnx.features import FeaturesManager

>>> distilbert_features = list(FeaturesManager.get_supported_features_for_model_type("distilbert").keys())
>>> print(distilbert_features)
["default", "masked-lm", "causal-lm", "sequence-classification", "token-classification", "question-answering"]

그런 다음 transformers.onnx 패키지의 --feature 인수에 이러한 기능 중 하나를 전달할 수 있습니다. 예를 들어, 텍스트 분류 모델을 내보내려면 다음과 같이 Hub에서 미세 조정된 모델을 선택하고 실행할 수 있습니다:

python -m transformers.onnx --model=distilbert-base-uncased-finetuned-sst-2-english \
                            --feature=sequence-classification onnx/

다음과 같은 로그가 표시됩니다:

Validating ONNX model...
        -[✓] ONNX model output names match reference model ({'logits'})
        - Validating ONNX Model output "logits":
                -[✓] (2, 2) matches (2, 2)
                -[✓] all values close (atol: 1e-05)
All good, model saved at: onnx/model.onnx

이때 미세 조정된 모델의 출력명은 이전의 distilbert-base-uncased 체크포인트에서 봤던 last_hidden_state와 달리 logits입니다. 시퀀스 분류를 위해 미세 조정되었기 때문에 예상대로 입니다.

지원되지 않는 아키텍처를 위한 모델 내보내기

라이브러리에서 직접 지원하지 않는 아키텍처의 모델을 내보내려면 세 가지 주요 단계를 거쳐야 합니다:

1. 사용자 정의 ONNX 구성을 구현하기
2. 모델을 ONNX로 내보내기
3. PyTorch 및 내보낸 모델의 출력 검증하기

이 섹션에서는 DistilBERT가 어떻게 구현되었는지 각 단계마다 자세히 살펴보겠습니다.

사용자 정의 ONNX 구성을 구현하기

ONNX 구성 객체부터 시작해 봅시다. 내보내려는 모델 아키텍처 유형에 따라 상속해야하는 세 가지 추상 클래스를 제공합니다:

• 인코더 기반 모델은 OnnxConfig를 상속합니다.
• 디코더 기반 모델은 OnnxConfigWithPast를 상속합니다.
• 인코더-디코더 모델은 OnnxSeq2SeqConfigWithPast를 상속합니다.

DistilBERT는 인코더 기반 모델이므로 해당 구성은 OnnxConfig를 상속합니다.

>>> from typing import Mapping, OrderedDict
>>> from transformers.onnx import OnnxConfig


>>> class DistilBertOnnxConfig(OnnxConfig):
...     @property
...     def inputs(self) -> Mapping[str, Mapping[int, str]]:
...         return OrderedDict(
...             [
...                 ("input_ids", {0: "batch", 1: "sequence"}),
...                 ("attention_mask", {0: "batch", 1: "sequence"}),
...             ]
...         )

각 구성 객체는 inputs 속성을 구현하고 매핑을 반환해야 합니다. 매핑의 키는 예상 입력에 해당하고 값은 해당 입력의 축을 나타냅니다. DistilBERT의 경우 input_ids 및 attention_mask 두 개의 입력이 필요한데요. 두 입력 모두 (batch_size, sequence_length)의 동일한 차원이기 때문에 구성에서도 똑같은 축을 사용합니다.

ONNX 구성을 구현한 후에는 다음과 같이 기본 모델의 구성을 제공하여 인스턴스화 할 수 있습니다:

>>> from transformers import AutoConfig

>>> config = AutoConfig.from_pretrained("distilbert-base-uncased")
>>> onnx_config = DistilBertOnnxConfig(config)

결과 객체에는 여러 가지 유용한 속성이 있습니다. 예를 들어 ONNX로 내보낼 때 쓰일 ONNX 연산자 집합을 볼 수 있습니다:

>>> print(onnx_config.default_onnx_opset)
11

다음과 같이 모델에 연결된 출력을 볼 수도 있습니다:

>>> print(onnx_config.outputs)
OrderedDict([("last_hidden_state", {0: "batch", 1: "sequence"})])

출력 속성이 입력과 동일한 구조임을 유의하세요. 각 출력은 이름과 차원이 OrderedDict의 키-값으로 저장되어 있습니다. 출력 구조는 구성을 초기화할 때 선택한 기능과 관련이 있습니다. 기본적으로 ONNX 구성은 AutoModel 클래스로 가져온 모델을 내보낼 때 쓰이는 default 기능으로 초기화됩니다. 다른 태스크를 위해 모델을 내보내려면 ONNX 구성을 초기화할 때 task 인수에 다른 기능을 넣으면 됩니다. 예를 들어, 시퀀스 분류 단계를 덧붙인 DistilBERT를 내보내려면, 이렇게 해볼 수 있습니다:

>>> from transformers import AutoConfig

>>> config = AutoConfig.from_pretrained("distilbert-base-uncased")
>>> onnx_config_for_seq_clf = DistilBertOnnxConfig(config, task="sequence-classification")
>>> print(onnx_config_for_seq_clf.outputs)
OrderedDict([('logits', {0: 'batch'})])

모델 내보내기

ONNX 구성을 구현했다면, 다음 단계는 모델을 내보내는 것입니다. 이제 transformers.onnx 패키지에서 제공하는 export() 함수를 살펴보겠습니다. 이 함수는 ONNX 구성, 기본 모델, 토크나이저, 그리고 내보낼 파일의 경로를 입력으로 받습니다:

>>> from pathlib import Path
>>> from transformers.onnx import export
>>> from transformers import AutoTokenizer, AutoModel

>>> onnx_path = Path("model.onnx")
>>> model_ckpt = "distilbert-base-uncased"
>>> base_model = AutoModel.from_pretrained(model_ckpt)
>>> tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

>>> onnx_inputs, onnx_outputs = export(tokenizer, base_model, onnx_config, onnx_config.default_onnx_opset, onnx_path)

export() 함수가 반환하는 onnx_inputs와 onnx_outputs는 구성의 inputs와 outputs 속성에서 정의된 키 목록입니다. 모델을 내보낸 후 다음과 같이 모델이 잘 구성되어 있는지 테스트할 수 있습니다:

>>> import onnx

>>> onnx_model = onnx.load("model.onnx")
>>> onnx.checker.check_model(onnx_model)

모델의 출력 검증하기

마지막 단계는 기존 모델과 내보낸 모델의 출력이 일정한 오차 범위 내에서 동일하다는 것을 검증하는 것입니다. 그러려면 transformers.onnx 패키지에서 제공하는 validate_model_outputs() 함수를 사용할 수 있습니다:

>>> from transformers.onnx import validate_model_outputs

>>> validate_model_outputs(
...     onnx_config, tokenizer, base_model, onnx_path, onnx_outputs, onnx_config.atol_for_validation
... )

이 함수는 generate_dummy_inputs() 메소드로 기존 및 내보낸 모델의 입력을 생성하며, 검증에 사용될 오차 범위는 구성에서 정의할 수 있습니다. 일반적으로는 1e-6에서 1e-4 범위 내에서 합의하지만, 1e-3보다 작다면 문제 없을 가능성이 높습니다.

🤗 Transformers에 새 구성 추가하기

미리 제작된 구성의 숫자를 늘리려고 노력하고 있으며, 커뮤니티의 기여를 환영합니다! 라이브러리에 당신만의 구성을 추가하려면 다음 단계를 기억해주세요:

• configuration_<model_name>.py 파일에 ONNX 구성을 구현하세요.
• ~onnx.features.FeatureManager에 모델 아키텍처 및 해당 기능을 포함하세요.
• test_onnx_v2.py의 테스트에 모델 아키텍처를 추가하세요.

아직 감이 안 잡히신다면, IBERT 구성이 어떻게 기여되었는지 확인해보세요.