Sooftware Serving - Huggingface Optimum

Soohwan Kim
Co-founder/A.I. engineer at TUNiB.
More posts by Soohwan Kim.

Soohwan Kim

05 Oct 2022•3 min read

Sooftware Serving - Huggingface Optimum

허깅페이스에서 나온 Transformers의 Extension 라이브러리이다. 목적은 모델 학습 및 인퍼런스를 더욱 빠르게 해주기 위한 라이브러리이다.

Exporting Transformers models to ONNX

모델 서빙을 위해서는 많이들 Transformers 모델을 ONNX로 컨버팅하곤 한다.
ONNX(Open Neural Network Exchange)는 쉽게 말하면, Tensorflow, PyTorch와 같이 서로 다른 딥러닝 프레임워크 환경에서 만들어진 모델들을 서로 호환해서 사용할 수 있도록 도와주는 오픈소스이다.

원래 Transformers 모델을 ONNX로 모델로 컨버팅하는게 마냥 간단하지만은 않았는데, 이 Optimum 라이브러리로 아래 코드로 쉽게 컨버팅 가능해졌다. 아마 Optimum을 가장 많이 쓰게되는 이유일 것 같다.

from optimum.onnxruntime import ORTModelForSequenceClassification

pretrain_model_name_or_path = "pretrain_model_name_or_path"
save_directory = "tmp/onnx/"

# Load a model from transformers and export it to ONNX
ort_model = ORTModelForSequenceClassification.from_pretrained(pretrain_model_name_or_path, from_transformers=True)

# Save the onnx model
ort_model.save_pretrained(save_directory)

Quantization

Model Quantization도 이 라이브러리로 쉽게 가능하다! Huggingface 만세!

from optimum.onnxruntime.configuration import AutoQuantizationConfig
from optimum.onnxruntime import ORTQuantizer

# Define the quantization methodology
qconfig = AutoQuantizationConfig.arm64(is_static=False, per_channel=False)
quantizer = ORTQuantizer.from_pretrained(ort_model)

# Apply dynamic quantization on the model
quantizer.quantize(save_dir=save_directory, quantization_config=qconfig)

Example of how to load an ONNX Runtime model and generate predictions with Optimum 🤗:

from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import pipeline, AutoTokenizer

model = ORTModelForSequenceClassification.from_pretrained(save_directory, file_name="model_quantized.onnx")
tokenizer = AutoTokenizer.from_pretrained(save_directory)

cls_pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)

results = cls_pipeline("I love burritos!")

Graph Optimization

Graph Optimization을 통해 모델 인퍼런스를 더욱 빠르게도 가능하다.

from optimum.onnxruntime import ORTModelForSequenceClassification, ORTOptimizer

ort_model = ORTModelForSequenceClassification.from_pretrained(pretrain_model_name_or_path, from_transformers=True)
optimizer = ORTOptimizer.from_pretrained(ort_model)

# Optimize the model
optimizer.optimize(save_dir=save_directory, optimization_config=optimization_config)