[ML OPS] quantization을 활용한 인퍼런스 최적화 (ft. ONNX, TensorRT)

오늘은 quantization을 활용하여 인퍼런스를 최적화하는 방법을 소개하려고 합니다.

이전 글(https://seokhyun2.tistory.com/82)을 같이 보시는 것을 추천드립니다.

 

0. quantization

quantization은 한국말로는 양자화라고 하는데요.

신호처리 수업을 들으신 분들은 이미 익숙하실 수도 있을 것 같습니다.

신호처리에서는 숫자들을 큰 숫자 셋에서 작은 숫자 셋으로 매핑하는 것을 양자화라고 합니다.

 

예를 들면, 1~1000을 10으로 나누고 정수만을 취해서 1~100으로 변경한다고 생각해보겠습니다.

511, 519가 동일하게 51이 되겠죠?

그러면 우리는 기존 데이터를 더 적은 숫자 셋(100개)으로 표현할 수 있지만, 511과 519가 동일한 데이터가 되면서 511과 519가 다르다는 정보는 없어지므로 데이터 자체에 대해서는 손실이 생기게됩니다.

신호처리에서는 압축을 위해서 사용하는 방식인데요.

quantization을 수행하게 되면, 완전히 손실이 되면 원상복구할 수 있는 방법이 없습니다.

 

quantization을 활용한 inferecne 최적화도 위의 원리와 동일합니다.

압축을 통해 데이터가 조금 손실이 되더라도, 메모리 사용량이 줄어들고 속도가 빨라지는 것으로 생각해주시면 됩니다.

 

딥러닝에서의 quantization에 대해서 조금 더 설명을 드리면, 학습된 weight들은 보통 소수로 저장이 되어있는데요.

기본적으로 weight를 저장하기 위해 FP32(Floating Point 32-bits)를 활용하게 됩니다.

컴퓨터는 실제로는 0과 1로 이루어진 이진수로 표현이 되는데, 32 bits를 활용하면 표현 범위가 2의 32제곱이 되겠죠.

FP32 대신에 FP16, INT8과 같은 타입을 활용하여, 손실이 생기지만 그만큼 속도가 빨라지게 할 수 있습니다.

 

참고로 FP16은 32에 비해서 bits를 절반만 쓰면서 정확도가 낮아지기 때문에 half precision이라고도 부르며, 최근에는 인퍼런스 뿐만 아니라 학습에서도 FP16만 활용하거나 FP16과 FP32를 섞어서 활용(mixed precision)하는 방식도 활용이 되고 있습니다.

학습 속도가 빨라지는 효과도 노릴 수 있고, FP32로 학습 후에 인퍼런스 시에 quantization을 따로 수행하면 손실이 생기면서 모델의 성능이 학습 결과로 기대했던 것보다 낮게 나오는 문제가 발생할 수 있지만, 처음부터 FP16으로 학습하고 인퍼런스도 동일하게 되면 손실이 생기지 않으므로 학습 결과를 그대로 서비스에 활용할 수 있다는 장점도 있습니다.

 

1. ONNX quantization

ONNX에서 제공하는 quantization은 8 bit quantization만 지원됩니다. https://onnxruntime.ai/docs/performance/quantization.html

Quantize ONNX Models

colab에서 bert 모델에 한번 적용을 해보도록 하겠습니다.

전체 코드는 여기(https://github.com/hsh2438/MLops/blob/main/5_quantization/onnx_quantization.ipynb)에서 확인하시면 됩니다.

 

ONNX를 활용하여 모델을 변환하는 코드는 간단합니다.

import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

model_fp32 = 'model/model.onnx'
model_quant = 'model/model.quantized.onnx'
quantize_dynamic(model_fp32, model_quant)

실행 후에 실제 용량을 한번 확인해보면 아래와 같이 32bit를 8bit로 줄였기 때문에, 용량이 4분의 1이 되는 것을 볼 수 있습니다.

그러면 이제 속도가 얼마나 빨라지는지 확인을 해봐야겠죠?

CPU에서 아래와 같이 간단하게 테스트를 해보면 속도는 빨라진 것을 확인할 수 있습니다.

하지만 32bit를 8bit로 줄이면 손실이 크기 때문에, quantization이 적용된 모델을 사용하실 때는 모델 성능에 문제가 없는지 꼭 확인 해보시는 것을 추천드립니다.

 

2. TensorRT

TensorRT(https://github.com/NVIDIA/TensorRT)는 NVIDIA에서 제공하고 있어서 NVIDIA GPU에서만 사용할 수 있습니다.

TensorRT는 colab에서는 사용이 너무 힘들기 때문에, 별도의 GPU 머신에서 테스트하였습니다.

 

위에서, ONNX로 quantization을 해보았는데 TensorRT는 왜 또 소개하는 것일까 궁금하텐데요.

ONNX는 8bit만 지원하는데, TensorRT는 16bit를 지원하기 때문입니다.

TensorRT에서는 FP16을 활용하면, INT8보다 성능 저하가 적으면서도 FP32에 비하여 inference 속도가 약 2배 빨라지게 됩니다.

 

TensorRT는 CUDA와 의존성 이슈도 있기 때문에, 설치도 어려운 편인데요.

하지만 친절하게도 이미 세팅이 다 되어 있는 docker image를 제공하기 때문에 편하게 받아서 활용해보도록 하겠습니다.

TensorRT가 설치된 docker image는 아래 링크에 접속해보시면, 확인해보실 수 있습니다.

https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tensorrt

TensorRT | NVIDIA NGC

 

참고로, docker에서 GPU를 활용하시려면, nvidia-docker도 설치되어 있어야 합니다.

 

이제 환경부터 세팅을 해보겠습니다.

아래 명령어로 docker container를 실행하면 image가 없는 경우에는 자동으로 다운 받아 와서 container가 실행이 됩니다.

docker run -it --gpus all nvcr.io/nvidia/tensorrt:22.06-py3 /bin/bash

 

GPU 관련 라이브러리와 python3이 이미 모두 설치되어 있는 이미지이므로, pip 명령어로 필요한 파이썬 패키지들을 설치해주시면 됩니다.

 

TensorRT를 사용하는 방법도 여러가지 방법이 있는데, onnxruntime에서도 TensorRT를 활용할 수 있도록 제공하고 있습니다.

자세한 내용은 아래 링크에서 확인하시면 됩니다.

https://onnxruntime.ai/docs/execution-providers/TensorRT-ExecutionProvider.html

TensorRT

 

그러므로, onnxruntime을 포함한 필요한 패키지들을 아래 명령어로 설치해주겠습니다.

pip install transformers torch onnxruntime-gpu

저희는 onnx로 변환된 bert 모델을 테스트하고 있으므로, 아래 명령어로 모델을 준비해주겠습니다.

python -m transformers.onnx --model=bert-base-uncased model

 

그러면 벌써 모든 준비는 끝이 났습니다. docker image에 이미 대부분 설치되어 있어서 정말 금방 구성할 수 있습니다.

이제 TensorRT가 얼마나 빠른지 한번 비교를 해보겠습니다.

 

환경만 잘 구성되었다면 onnxruntime에서 TensorRT를 사용하는 방법은 간단합니다.

CUDAExecutionProvider 대신에 TensorrtExecutionProvider를 사용하시면 됩니다.

모델을 따로 변환하는 과정은 필요없지만, 대신에 모델을 로딩하는데 시간이 조금 더 소요됩니다.

 

추가로 FP16을 활용하기 위해서는, trt_fp16_enable 옵션을 True로 설정해주시면 됩니다.

 

전체 코드는 아래 링크에서 확인하시면 됩니다.

https://github.com/hsh2438/MLops/blob/main/5_quantization/trt_fp16.py

 

위의 코드를 실행하면, 아래와 같은 결과를 얻을 수 있는데요.

TensorRT 사용하는 것 만으로도, 속도가 개선이 되고 FP16 옵션까지 활용했을땐 2배 이상 빨라지는 것을 확인할 수 있습니다.

 

FP16도 어쨌든 손실은 있으므로, 모델에 따라서 성능 저하가 많이 일어날 수 도 있으므로 서비스에 활용하시기 전에 꼭 성능 저하가 없는지 확인하시는 것을 추천드립니다.