[Project] wav2vec2 모델 ONNX+TritonInferenceServer로 배포하기

개발/Project 2024. 4. 15. 22:03

Huggingface 는 최신 논문부터 사람들이 직접 학습한 모델까지 다양한 딥러닝 모델을 사용할 수 있는 플랫폼이다. 하지만 실제로 python 환경에서 huggingface 모델들을 사용해보면 model loading, inference latency 등 생각보다 많은 문제가 있어서 huggingface 모델 자체로 프로덕트 배포하기에는 어려움이 있다. 이런 문제들을 해결하기 위해 다양한 ML 프레임워크가 존재하는데 이번 글에서는 wav2vec2 모델을 ONNX 로 변환한 후 TritonInferenceServer 를 통해 배포하는 과정을 소개한다. 추가적인 스크립트 작성없이 몇가지 커맨드만으로 재현할 수 있도록 만들어놨다

왜 ONNX + TritonInferenceServer ?

ML 모델은 pytorch, tensorflow 등 다양한 ML 프레임워크에서 학습이 진행되는데 각 프레임워크간에 호환성 문제가 존재한다. 그 문제를 해결하기 위해서 학습한 딥러닝 모델을 ONNX (Open Neural Network Exchange) 로 배포하여 어느 ML 프레임워크에서 학습을 하던간에 배포를 용이하게 하려는 목적이 있다.

그렇다면 TritonInferenServer 는 왜 사용할까? 이건 서비스의 특징 (or 딥러닝 모델에서 보완하고 싶은 문제) 에 따라 달라지는데, 내가 다니던 회사는 실시간 음성인식을 서비스하는 회사여서 인공지능 모델로 들어오는 다양한 길이의 오디오를 비동기 디코딩하여 빠르게 처리해야했다. Nvidia 의 TritonInferenceServer 는 이런 문제를 해결하기에 적합한 툴이였으며, gRPC 를 통한 데이터 통신 또한 편리해서 선택했다.

1. ONNX 모델 export

들어가기에 앞서 실험환경은 다음과 같다. 많이 열악하지만 이런 적은 리소스만으로도 충분히 재현가능하다!

nvidia driver version: 535.154.05
GPU: GTX 1080 VRAM 8GB
RAM: 16GB
Ubuntu 22.04
Nvidia-Docker (GPU available)

실제로 ONNX 로 변환할 모델은 huggingface 에서 wav2vec2 한국어 모델 중 가장 다운로드를 많이 받은 kresnik/wav2vec2-large-xlsr-korean 모델이다. 어떻게 학습됐는지 살펴보니 아래와 같이 한국어 음절(syllable) 단위로 학습되었다. 총 음절수는 1,203 개로 한국어로 표현할 수 있는 음절 수에는 한참 못미치지만, 이번 글의 목적은 모델의 인식률이 아니므로 넘어가도록 하자. 만약 자모 (grapheme) 단위로 학습된 한국어 모델을 찾는다면 thisisHJLee/wav2vec2-large-xls-r-300m-korean-g 모델도 있으니 참조하자.

{'볍': 0, '칠': 1, '깊': 2, '뭔': 3, '러': 4, '르': 5, '튀': 6, '쳇': 7, '땀': 8, '픔': 9, '밌': 10, '좁': 11, '찧': 12, '뮬': 13, '했': 14, '연': 15, '핵': 16, '붐': 17, '봇': 18, '궁': 19, '뜸': 20, '넌': 21, '젖': 22, '맏': 23, '벚': 24, '락': 25, '가': 26, '롭': 27, '달': 28, '슐': 29, '컸': 30, '읍': 31, '색': 32, '맛': 33, '닦': 34, '기': 35, '악': 36, '뻗': 37, '팅': 38, '끗': 39, '깝': 40, '후': 41, '소': 42, '끽': 43, '조': 44, '겪': 45, '코': 46, '광': 47, '컨': 48, '올': 49, '큰': 50, '델': 51, '굉': 52, '교': 53, '발': 54, '랍': 55, '난': 56, '맹': 57, '킨': 58, '옵': 59, '삭': 60, '녹': 61, '엄': 62, '마': 63, '쫄': 64, '콜': 65, '넥': 66, '웹': 67, '뻔': 68, '뿐': 69, '엠': 70, '태': 71, '헝': 72, '멕': 73, '괭': 74, '멧': 75, '핀': 76, '냉': 77, '덟': 78, '덕': 79, '딪': 80, '넷': 81, '향': 82, '왜': 83 ...

환경셋업

평소 Docker 를 사용하는 것을 선호해서 모든 환경 세팅은 Docker 기반이다. 아래 커맨드를 이용해서 ONNX 변환을 위한 컨테이너를 띄어주자.

git clone https://github.com/coolseaweed/wav2vec-triton.git
cd wav2vec-triton && docker compose --profile export up -d --build

ONNX 모델 변환

이제 아래 커맨드를 이용해서 huggingface wav2vec2 모델을 ONNX로 변환해보자. 변환하는 스크립트는 run.py 를 통해 확인 할 수 있다. 만약 FP16 으로 변환하고 싶다면 --fp16 파라미터를 명시해주자. 아래에도 비슷한 커맨드들이 나타나는데 FP16 관련 설명은 추가로 더하진 않겠다.

docker exec wav2vec2-triton-export-1 python run.py [--fp16]

성공적으로 실행되었다면 실제로 ONNX 모델을 inference 해보자.

docker exec wav2vec2-triton-export-1 python inference.py [--fp16]

아래와 같이 디코딩 텍스트가 나온다면 잘 변환 되었음을 알 수 있다.

decoding result: 이 일 자동차업계에 따르면 임팔라는 지난달 삼십 일 일 하루에만 구백 대가 넘는 계약이 이뤄진 것으로 파악됐다

아래와 같은 파일구조가 생성되었다면 다음으로 넘어가자.

models/
├── config.pbtxt
└── wav2vec2
    ├── 1
    │   └── model.onnx
    └── config.pbtxt

2. Triton Inference Server 세팅

환경세팅

Triton Inference Server 를 이용한 inference를 위해서는 server 와 client 두가지 환경세팅이 필요하다.

docker compose --profile server --profile client up -d --build

docker-compose.yaml 파일을 참조해보면 localhost 간에 아래와 같이 port forwarding 이 되어있음을 알 수 있다.

  server:
    <<: *base
    build: ./triton
    ports:
      - "8000:8000"
      - "8001:8001"
      - "8002:8002"

    profiles: ['server','all']
    volumes:
      - *v_models
    # command: 'tritonserver --model-repository /models --log-verbose 1' # debug

localhost의 8000번 포트를 이용해서 triton server 가 잘 떠있는지 확인해보자

curl -v localhost:8000/v2/health/ready
----------------------------------
*   Trying 127.0.0.1:8000...
* Connected to localhost (127.0.0.1) port 8000 (#0)
> GET /v2/health/ready HTTP/1.1
> Host: localhost:8000
> User-Agent: curl/7.81.0
> Accept: */*
> 
* Mark bundle as not supporting multiuse
< HTTP/1.1 200 OK
< Content-Length: 0
< Content-Type: text/plain
< 
* Connection #0 to host localhost left intact

마지막으로 client 에서 triton server 를 통해 inference 가 잘 되는지 확인해보자. 다른 audio도 test 해보고 싶다면 data 디렉토리에 auido file을 넣고 -i 커맨드로 경로를 설정해주자.

docker exec wav2vec2-triton-client-1 python inference.py -i test_audio.wav [--fp16]

아래와 같이 디코딩텍스트가 나온다면 성공

이|일|자동차업계에|따르면|임팔라는|지난달|삼십|일|일|하루에만|구백|대가|넘는|계약이|이뤄진|것으로|파악됐다

여기까지 잘 따라왔다면 이제 onnx 로 변환된 wav2vec2 모델을 triton inference server 에서 동작시킬 준비가 완료되었다.

3. Benchmark

실제로 ONNX 로 변환된 모델이 TritonInferenceServer 에서 동작하면서 개선된점이 어떤게 있는지 아래 표를 보면서 확인해보자. 벤치마킹에 사용된 데이터셋은 오픈소스인 zeroth 테스트셋을 사용했으며, 오디오 총 길이는 4287 초 분량이다.

huggingfaceONNX(FP32) + TritonONNX(FP16)+Triton

	HuggingFace	ONNX(FP32)+Triton	ONNX(FP16)+Triton
WER (%)	4.74	4.73	4.70
CER (%)	1.78	1.78	1.77
processing time (sec)	49.96	41.53	35.98
RTF	0.0115	0.0097	0.0084
GPU (MiB)	3916	3286	1744

우선 baseline 인 huggingface 모델과 ONNX (FP32) + Triton 간의 인식률을 나타내는 WER(Word Error Rate) 와 CER(Character Error Rate) 을 비교해보면 성능감소는 없었다. 오디오를 디코딩하는데 드는 총 시간 (processing time) 이 ONNX (FP32) + Triton 에서 약 20% 빨라진걸 확인 할 수 있었으며, GPU 사용량도 17% 감소한걸 확인할 수 있었다. ONNX (FP16) + Triton 모델은 half precision 임에도 불구하고 이상하게도 인식률이 baseline 보다 좋게 나왔다. 테스트셋이 1시간 분량밖에 되지 않기 때문에 인식률의 오차가 있어보인다. 좀더 큰 테스트셋으로 측정해보면 분명 인식률의 손실이 있을 것 같다. 반면에 processing time 은 약 38% 빨라졌으며, GPU 메모리도 절반이하로 줄어든 것을 확인할 수 있다. FP16 모델은 baseline 인식율을 떨어뜨릴 수도 있지만 컴퓨팅 리소스 절감이 뚜렸해서 trade off 를 잘 따져보면 좋을 것 같다.

벤치마킹 결과는 GPU 모델에 따라 다르게 나타날 수 있다. 벤치마킹을 시뮬레이션해보고 싶다면 아래 커맨드를 참조하자.

docker exec -it wav2vec2-triton-export-1 python benchmark.py # huggingface benchmark
docker exec -it wav2vec2-triton-client-1 python benchmark.py [--fp16] # triton benchmark

마치며

언젠가 블로그로 정리하고 싶었던 경험 중 하나였다. 최대한 쉽게 따라해 볼수 있도록 코드를 정리하다보니 시간이 많이 들었지만, 이렇게 마침표 하나를 찍을 수 있어서 마음이 편안하다. 실제로 ONNX + TritonInferenceServer 는 현재까지 (2024.01.25) 많이 사용되는 툴이기에 널리 알리고 많은 사람들이 잘썼으면 하는 바람에 글을 남긴다. 물론 내가 놓치고 있는 부분이나 추가적으로 개선해야할 부분들을 코멘트로 남겨주면 감사하겠다. 코드나 잘 안되는 부분은 댓글 또는 github 이슈로 남겨주면 답장드릴테니 많은 관심 바란다

ABOUT ME

식은김 개발 블로그