Optimum

이번 포스팅에서는 Optimum을 사용하여 huggingface hub에 올라와 있는 wav2vec2 모델을 onnx로 변환하는 내용을 다루고 있습니다.

huggingface 모델

huggingface hub에 가보시면 다양한 언어로 학습한 whisper, wav2vec2 모델들이 올라와 있는 걸 확인하실 수 있습니다.

많은 모델들이 있지만 예제로는 meta에서 공개한 wav2vec2-base-960h 모델을 사용하여 진행하겠습니다.

Onnx

Optimum으로 onnx 모델 변환을 하기 전에 onnx에 대해서 간략하게 알아보겠습니다.

ONNX(Open Neural Network Exchange)는 머신러닝 모델들을 다양한 프레임워크로 호환 가능하게 만들어주는 개방형 포맷입니다.
이 표준을 사용함으로써 하나의 프레임워크에서 훈련된 모델을 다른 프레임워크로 쉽게 이동시켜 사용할 수 있습니다.
예를 들어, PyTorch에서 훈련된 모델을 onnx로 변환하여 TensorFlow에서 실행할 수 있습니다.

onnx는 Microsoft, Amazon, Facebook 등 여러 기업이 공동으로 개발했으며, 여러 머신러닝과 딥러닝 프레임워크들의 지원을 받고 있습니다. 또한 onnx의 중요한 특징으로는 호환성, 확장성이 좋다고 할 수 있습니다.

PyTorch 모델을 onnx 모델로 변환하려면 torch.onnx.export 함수를 사용해야 되는데요.
Optimum을 사용하면 huggingface 모델을 load해서 내부적으로 torch.onnx.export 까지 작동하기 떄문에 쉽게 onnx 모델을 변환할 수 있습니다.

onnx 모델을 변환하기 전에 TorchScript로 변환해야 되는데요.

자세한 내용은 해당 포스트에서 확인해주세요.

Optimum을 사용하여 모델 변환

Optimum을 설치하시면 정말 간단하게 모델을 변환하실 수 있습니다.

MODEL_NAME=facebook/wav2vec2-base-960h
SAVE_PATH=/path/to/save
OPSET=15
python -m optimum.exporters.onnx --model "${MODEL_NAME}" "${SAVE_PATH}" --task automatic-speech-recognition --opset "${OPSET}" --framework pt