brunch

You can make anything
by writing

C.S.Lewis

by Qscar Mar 02. 2023

[AI]Optimization Technique 04

onnx

4. Distillation + Quantization



복습

지난 시간까지 Quantization, Pruning, Distillation에 대해 알아보았습니다. 이번 포스팅에서는 이러한 최적화 기법 중 Quantization과 Distillation을 함께 사용함으로써 original 대비 얼마나 최적화가 가능한지, 추가적인 최적화 방법은 없는지 살펴보도록 하겠습니다.


우선 지난 세 번째 포스팅 과정에서 다뤘던 모델들을 불러오도록 하겠습니다. 정확히는 Original Bert 모델을 clinc 데이터셋(의도분류)에 대해 파인튜닝한 것과 이 모델을 지식정제한 모델입니다. 다음의 코드를 통해 Original Bert는 Huggingface Hub에서, 지식정제한 모델은 로컬에서 가져옵니다.



Model Load

Load Original Bert & Quantized Bert - Image by author

이렇게 불러온 코드를 transformers 라이브러리를 이용해  text-classification task에 대한 pipeline을 구축합니다. 각각 original과 distil이라는 이름으로 정의해준 뒤, benchmark를 확인해보도록 하겠습니다. 사용하는 benchmark는 최적화 기법의 첫 번째와 세 번째 포스팅에서 사용한 것과 동일한 클래스입니다. 평가 데이터셋으로는 학습 과정에서 사용되지 않은 clinc dataset의 test셋을 사용합니다.



Original Model

Original Bert Benchmark - Image by author

우선 지난 시간에도 살펴봤던 original bert에 대한 benchmark 결과입니다. Size는 418mb, 평균 실행속도는 15.25ms, 정확도 성능은 0.8673입니다. 참고로 평균 실행속도는 실행 시마다 다소의 차이가 있을 수 있으며 대략 20% 내외의 오차를 보이곤 합니다.



Distilated Model

Distil Bert Benchmark - Image by author

지식정제한 모델의 결과입니다. 사이즈는 256mb, 평균 실행속도는 8.09ms, 정확도 성능은 0.8705입니다. 모든 측면에서 원본 모델보다 뛰어난 결과입니다.


여기서 우리는 크게 두 가지를 고민할 수 있습니다.

1. 원본 모델을 양자화하는 것
2. 지식정제한 모델을 양자화하는 것

이에 대해 모두 확인해보도록 하겠습니다.

첫 번째 포스팅을 통해 우리는 torch 라이브러리의 quantize_dynamic 메서드를 통해 간단히 양자화할 수 있음을 확인했습니다. 이번 시간에도 이를 동일하게 응용해 비교해보도록 하겠습니다.



Quantized Model

Quantized Bert Benchmark - Image by author

original model을 바로 양자화한 결과입니다. 모델 사이즈는 173mb, 평균 실행속도는 13.93ms, 정확도 성능은 0.8635입니다. 모델의 사이즈는 여태까지 진행한 모든 기법 중 제일 작지만 실행속도는 지식정제 모델에 비해 느리며(원본 모델보단 빠르지만), 정확도 성능은 가장 낮습니다.



Distil & Quanited Model

Distillation + Quantization Bert Benchmark - Image by author

다음으로는 지식정제한 모델을 양자화한 결과입니다. 모델 사이즈는 132mb, 평균 실행속도는 6.7ms, 정확도 성능은 0.8665입니다. 모델 사이즈는 가장 작고, 평균 실행속도 또한 가장 빠르며, 정확도 성능은 원본 모델보단 살짝 낮지만 이 정도면 거의 같은 수준입니다.


위 모델들의 벤치마크 결과를 정리하면 다음과 같습니다.

최적화 기법 별 Benchmark 정리 - Image by author

위 표에서 수치 옆의 괄호는 순위를 의미합니다. 만약 성능을 떨어뜨리지 않는 것이 중요하다면 지식정제 모델을, 메모리 최적화 등을 더 우선으로 둔다면 지식정제 모델을 양자화한 것을 선택하는 것이 합리적으로 보입니다.






5. ONNX


What is ONNX?

ONNX는 이전에 언급했던 torch, tensorflow와 같이 하나의 라이브러리입니다. 정확히는 딥러닝 모델을 나타내기 위한 공통 연산자와 공통 파일 포맷을 정의하는 공개표준입니다. ONNX 라이브러리를 통해 ONNX 포맷으로 변환이 가능해지며, ONNX 포맷으로 변환된 모델은 중간 표현(Intermdediate Representation)이라 불리는 계산 그래프를 만들게 됩니다. 이 라이브러리를 통해 torch로 구축한 모델을 tensorflow로 변환하거나, 그 반대, 혹은 그 외의 변환도 가능해집니다.


하지만 일반적으로 ONNX는 ONNX Runtime(보통 'ORT'로 줄여서 씀)와 함께 사용할 경우, 사용하는 하드웨어 종류에 걸맞은 성능 최적화를 기대할 수 있습니다. 그 이유는 ONNX 포맷으로 변환된 모델에 대해 연산자 융합(operator fusion), 상수폴딩(constant foliding)과 같이 ONNX 그래프 연산을 최적화시키기 때문입니다. 또한 실행 공급자(Execution Provider)를 통해 하드웨어에 대한 인터페이스를 정의합니다.


이러한 ONNX 변환 또한 대부분의 라이브러리에서 간단히 사용할 수 있도록 구현되어 있습니다. 여기에서는 huggingface의 convert_graph_to_onnx.convert() 내장함수를 통해 구현해보도록 하겠습니다.



ONNX Process

ONNX 포맷 변환과 이후의 작업을 위해서는 추가적인 전후 작업이 필요해지게 됩니다. 하나씩 정리해가며 진행하도록 하겠습니다.


1) 환경변수 지정

환경변수 지정 - Image by author

ONNX 포맷으로 변환하고, 성능을 최적화하기 위해서는 위 그림과 같이 환경변수에 대한 지정이 필요합니다. 병렬 작업을 위해 OMP API를 활성화하고(row 4), CPU 프로세서의 사이클을 사용하기 위해 대기 스레드를 활성 상태로 지정합니다(row 5).



2) ONNX 포멧 변환

convert onnx format - Image by author

ONNX 포맷으로 변환하기 위해 transformers 라이브러리의 convert_graph_to_onnx 내장함수를 사용합니다(row 1). pathlib 모듈의 Path 내장함수를 통해 변환한 ONNX 포맷의 모델을 저장할 경로를 지정하고(row 2~4), ONNX 포맷으로 변환하는 작업을 진행합니다. 참고로 ONNX 포멧 변환은 텐서타입의 모델만을 지원하기 때문에 양자화된 모델은 변환이 불가능합니다. 여기선 지식정제 모델을 사용하도록 하겠습니다.


ONNX 포맷으로 변환해주기 위해선 현재 모델이 어떤 프레임워크 기반 모델인지(row5), 그 모델은 어디에 있는지(row6), task에 따른 전처리 등을 위한 사전작업 모델은 없는지(row7, row10), 출력결과를 어디에 저장할 것인지(row8), onnx 라이브러리의 어떤 버전을 사용할 것인지(row9)를 지정합니다. 일반적으로 onnx 라이브러리 버전은 최신 버전을 사용할수록 더 빠르고, 성능이 좋다고 합니다. 이 글을 작성하는 2023년 03월 02일 기준으로는 17 버전이 가장 최신 버전입니다.



3) ONNX RUNTIME(ORT) 구현

Define ONNX Provider & Session class - Image by author

공식문서에 따르면, ONNX 포맷(엄밀히 말하면 'ONNX Graph Format')의 모델은 각기 다른 하드웨어에 최적화된 오퍼레이터 커널이 적용된 그래프 연산 방식으로 구성되어 있습니다. 조금 어려운 말인데요, 간단히 말하면 오퍼레이터처럼 작동합니다.


이를 위해선 operator kernels을 실행시킬 execution provider가 필요하고, 이 provider는 실행되는 하드웨어 타입에 맞는 kernel set을 구축합니다. 이때 provider가 지원하는 하드웨어는 CPU, GPU, IoT 등 다양하며, 지원하는 execution provider들은 여기서 확인할 수 있습니다. 


Supporting Optimization Functions as Provder - Image by onnxruntime docs[2]

또한 위 그림에서 보다시피 provider의 종류에 따라 ONNX 포맷의 최적화 기능에 차이가 있습니다. cpu는 GELU Approximation을 제외한 모든 최적화를 지원합니다. options 및 session 이하에는 profiling부터 복수의 provider들을 지정하는 등 각종 기능을 지원합니다.



4) ONNX PIPELINE 구현

onnx pipeline - image  by author

기존에는 transformers 라이브러리의 pipeline 모듈을 이용해 text-classification task에 대한 파이프라인을 자동으로 구축했습니다. 하지만 당연하게도 해당 기능은 onnx와 호환되지 않기 때문에 해당 모듈의 필요한 기능만을 위 코드를 통해 구현했습니다.


간단히 설명하면 query가 입력되면 tokenizer로 이를 토큰화하고(row 9), 이를 key, value로 나눠 onnx model에 넣습니다(row 10~11). 이후 전체 logit에 대해 softmax를 취하고, argmax 함수를 통해 가장 높은 확률의 인덱스를 확인합니다(row 12~13). 마지막으로 이러한 결과를 레이블과 확률값으로 전달합니다(row 15).



5) ONNX Benchmark 구현

onnx benchmark

Benchmark 또한 조금 변경을 해줘야 합니다. 기존에는 torch base model이었기 때문에, torch.save 메서드를 이용해 모델을 저장한 다음 모델 사이즈를 읽어들이고, 삭제하는 과정을 거쳤습니다. 하지만 여기선 이미 모델이 저장되어 있기 때문에 모델의 경로에 따른 사이즈를 측정하도록 업데이트해주면 됩니다. 


이를 위해 이전에 작성했던 PerformanceBenchmark 클래스를 승계해, size를 측정하는 함수만 업데이트해줍니다.



6) Benchmark

onnx model benchmark - Image by author

마지막으로 onnx graph format으로 변환한 모델의 성능을 확인해보겠습니다. 해당 변환은 지식정제 모델을 기반으로 했습니다. onnx 변환을 통해 모델 사이즈와 성능에는 변함이 없지만, 실행속도는 8.09ms에서 4.41ms로 두 배 가량 빨라졌음을 확인할 수 있습니다.


이러한 추론속도는 하드웨어의 성능, 해당 하드웨어에 대한 provider와 그에 따른 최적화 기능, 세부 option 등에 따라 달라질 수 있는데요. 만일 위 코드 중 ORT를 구축하는 단계에서 사용할 thread의 숫자를 변경하거나, provider를 gpu가 호환되도록 했을 경우 더 빠른 추론이 가능해집니다. GPU를 이용한 성능은 본 포스팅의 마지막을 통해 일괄적으로 확인해보도록 하겠습니다.



ONNX Format + Quantization


1) Quantization

ORT를 이용한 동적 양자화 - Image by author

ONNX 또한 내재 함수를 통한 양자화가 가능합니다. 여기선 이전에 했던 것과 동일하게 동적 양자화를 적용해보도록 하겠습니다. torch에서 했던 것과 거의 동일합니다.



2) Benchmark

Distil → ORT → Quantized Benchmark - Image by author

이렇게 양자화한 모델의 벤치마크를 확인해보면 위와 같습니다. 엄밀히 말하면 우리는 이 단계까지 오기 위해 Original Bert Model을 Distillation한 뒤, 해당 모델을 ORT(ONNX RunTime)를 이용해 ONNX Graph Format으로 변환했고, 다시 ORT를 이용해 양자화를 진행한 결과입니다. 


정확도 성능은 다소 떨어졌지만 모델 사이즈가 255mb에서 64mb로 무려 1/4로 줄어들었습니다!

nn.Linear 모듈 위주로 최적화하는 torch와 다르게 ONNX는 임베딩 층도 양자화하기 때문입니다(torch에서 nn.Linear 모듈 위주로 최적화하는 이유는 실제 대부분의 리소스가 fully connected layer에서 사용되는 nn.Linear에 집중되기 때문-VGGNet 기준 전체 파라미터의 80% 이상이 FC의 nn.Linear에 사용). 또한 실행속도 또한 3.8ms로 본래의 8.09ms에 비하면 배가 넘게 빨라진 것을 확인할 수 있습니다.



3) Summary

CPU 환경에서 최적화 기법 정리 - Image by author

결과적으로 지식정제한 모델을 ONNX RunTime을 적용해, 양자화한 모델이 대부분의 지표에 있어 가장 우수했습니다. 정확도 성능 또한 2위이긴 하지만, 본래 모델의 성능이었던 0.8673보다도 높은 성능을 보이고 있기에 만족할만한 수준입니다.



Running on GPU


GPU SETTING(CUDA)

마지막으로 간단히 GPU 환경에서 실행시키는 방법과 그 결과에 대해 정리하고 넘어가도록 하겠습니다.

우선 기존의 transformers 라이브러리를 통해 pipeline을 구축했다면 다음과 같이 pipeline 메서드를 사용할 때, device 매개변수를 추가해주면 됩니다. 이때 device 넘버에 따라 0 이상의 숫자를 적으면 되며, -1은 cpu, 복수의 gpu를 사용할 땐 별도의 추가작업이 필요합니다.

 

transformers에서 gpu 세팅 - Image by author

ORT에서는 Provider를 설정해주는 것으로 GPU를 할당해줄 수 있습니다. 다만 그 전에 'onnxruntime-gpu' 라이브러리를 추가로 설치해야합니다.


gpu 사용을 위한 provider 설정 - Image by author

위 그림과 같이 간단하게 세팅해주는 방법은 일반적으로 할당된 자원을 최대한 사용하는 식으로 작동합니다. 하지만 현실적으로는 다른 모델들이 동작하거나, 서비스 환경 유지 및 서비스 품질 등을 위해서라도 사용하는 컴퓨팅 리소스를 제한적으로 사용해야하는데요, 이때는 아래와 같은 코드를 통해 자원을 제약적으로 사용할 수도 있습니다.


Provider GPU 세팅 - Image by author

참고로 양자화된 모델에 대해선 gpu 사용을 지원하지 않습니다. 이는 양자화라는 개념 자체가 제한된 머신 리소스 환경에 사용할 것을 가정하고 만들어졌기 때문인데요, ORT를 이용하면 양자화한 모델을 굳이 GPU환경에서 돌릴 수 있지만, ORT의 최대 강점인 device에 대한 최적화가 이뤄지지 않기 때문에 오히려 속도가 더 느려지게 됩니다.



Summary

위 gpu 응용코드들을 통해 제 gpu 환경인 NVIDIA A100에서 동작시킨 결과는 다음과 같습니다. 일반적으로 GPU 머신 상에 돌리는 것으로 모델 사이즈나 정확도와 같은 성능지표의 차이는 발생하지 않기에(ORT로 양자화시킨 모델을 돌릴 때는 성능차이가 발생하기도 합니다), 기존에 정리했던 표에 G-TIME이라는 칼럼을 새로 만들어 정리해보도록 하겠습니다.


optimization 기법에 대한 정리 - Image by author

위 그림을 통해 알 수 있듯, GPU를 사용한다면 지식정제한 모델을 ORT로 변환한 모델이 무려 1.5ms 미만의 속도로 CPU환경에서 돌렸던 original 모델 대비 약 289배나 빠르게 동작합니다! 또한 ORT로 양자화한 모델은 GPU 환경에서 동작시켰을 경우, 속도와 정확도 성능이 모두 하락하는 결과를 확인할 수 있었습니다.


때문에 지금까지의 결과로 엣지 디바이스와 같이 소형 디바이스 혹은 머신에서는 양자화한 모델이 가장 선호될 것이며, 자원이 충분하다면 양자화를 하지 않은 모델이 가장 선호될 것임을 확인할 수 있었습니다.


여기까지 기본적인 최적화 방법에 대해 다뤄보았습니다.

다음 시간부터는 transformer 구조에 최적화된 transformers 라이브러리 내부에 구현된 xformer, accelerater와 같은 최적화 기능 등을 generation task에 적용하는 과정을 다뤄보도록 하겠습니다.



REFERENCE

[1] https://huggingface.co/transformers/v3.3.1/serialization.html#:~:text=Change%20the%20target%20opset%20version%20of%20the%20generated%20model.%20(More%20recent%20opset%20generally%20supports%20more%20operators%20and%20enables%20faster%20inference)

[2] https://onnxruntime.ai/docs/

브런치는 최신 브라우저에 최적화 되어있습니다. IE chrome safari