AI 경량화를 툴킷 하나로, OwLite -2/2
OwLite는 AI 모델을 쉽게 경량화하면서도 강력한 성능을 유지할 수 있는 툴킷입니다.
AI가 점점 다양한 서비스를 통해 우리 생활에 가까워지고 있습니다. 이와 함께 AI 기술 기반의 서비스를 지속가능한 형태의 비즈니스로 만들기 위해 AI 경량화 기술에 대한 필요성이 증가하고 있습니다.
이와 관련된 자세한 내용이 궁금하시다면 1편을 참고해주세요!
새롭게 생겨나는 경량화 수요에 대응하기 위해 OwLite는 이전에는 없던 많은 기술적, 기획적 고민을 거쳐 탄생했습니다. 오늘은 OwLite의 출시를 맞아 저희가 걸어온 과정을 소개해 보겠습니다. OwLite가 탄생한 과정과 고민을 알고 나면 OwLite 가 어떤 기능을 가지고 있는지, 여러분들에게 어떤 도움을 드릴 수 있는지 명확해질 것이라고 생각합니다.
기술
AI 모델을 ‘잘’ 경량화하기 위해서는 모델의 크기를 줄이고 추론 속도를 개선하는 동시에 모델의 성능이 저하되지 않도록 주의해야 합니다. 이를 위해 모델 구조에 대한 깊은 이해가 필수적이며, 필요에 따라 적절한 경량화 방법론을 적용하는 것이 중요합니다.
예를 들어, 정확도에 영향을 크게 미치지 않으면서도 Latency 감소를 크게 이룰 수 있는 Layer를 파악해 부분적으로 경량화 한다거나, 특정 위치나 종류의 Layer에 성능 저하를 줄일 수 있는 다른 경량화 알고리즘을 적용하는 등의 방법이 가능합니다. 하드웨어의 특성에 따라 경량화 방식을 다양화하는 것 역시 시도해 볼 수 있을 것입니다.
엔지니어의, 엔지니어에 의한, 엔지니어를 위한 도구
스퀴즈비츠 팀에서 실제로 여러 고객사들의 모델을 경량화하는 과정에서 반복되는 작업과 공통적으로 적용될 수 있는 기술들이 있음을 확인했고, OwLite는 스퀴즈비츠의 경량화 엔지니어가 더 효율적으로 모델을 경량화할 수 있도록 돕는 도구라는 컨셉으로 시작되었습니다.
경량화 엔지니어가 직접 사용하기 위해 만든 도구이기 때문에, OwLite는 경량화 과정의 많은 조건과 제약을 고려한 방식으로 기획되었습니다.
OwLite는 AI 모델 트레이닝에 널리 사용되는 PyTorch 프레임워크와 호환되는 Python 패키지, 그리고 AI 모델의 구조를 시각적으로 확인하고 이를 기반으로 경량화를 적용할 수 있는 GUI 로 이루어져 있습니다.
다재다능한 호환성
AI 모델을 시각화하고, 시각화 한 모델 상에서 경량화 방법론을 선택하고, 이를 실제 경량화 기술이 반영된 TensorRT로 만들어내는 과정은 Pytorch ↔ ONNX ↔ TensorRT간의 유기적인 호환을 필요로 합니다. 기존에도 각 프레임워크를 넘나들 수 있도록 하는 기술은 존재했지만, (1)모델을 시각화 하고 (2)시각화 된 모델에서 경량화를 적용하고 (3)이를 실제 경량화 된 모델에 반영하는 과정은 그보다 더 복잡한 기술을 필요로 합니다.
PyTorch로 학습 및 개발된 AI 모델은 ONNX와 같은 중간 단계를 거쳐 TensorRT 등의 inference 전용 프레임워크로 변환되어 서비스됩니다. 이 과정에는 모델을 효율적으로 구동할 수 있도록 돕는 컴파일 과정이 포함되어 있어 PyTorch → ONNX → TensorRT 의 방향으로만 변환이 가능하고 역방향으로의 변환이 불가능합니다. 그렇기 때문에 TensorRT 수준에서의 deployment를 고려한 채로 PyTorch 모델을 수정하거나 추가 학습하는 것이 매우 어려웠습니다.
스퀴즈비츠의 엔지니어들은 PyTorch 모델 그래프, ONNX 모델 그래프, 그리고 TensorRT 그래프 사이의 변환을 자유롭게 할 수 있는 컴파일러를 개발했고, 각 모델 간의 여러 layer들을 매칭하는 기술을 개발했습니다. 이를 통해 더 작은 단위에서 호환성을 가질 수 있도록 했습니다. 이는 후술할 다양한 이점을 만들어 내게 됩니다.
추가적으로, OwLite 는 엣지 디바이스 배포 등 더 넓은 니즈에 대응하고자 tflite와 같은 다양한 Framework 로의 확장을 계획하고 있습니다.
간편하지만, 더욱 더 강력한 경량화
TensorRT 프레임워크는 자체적으로 엔진 빌드 과정에서 기초적인 수준의 경량화 기술을 제공하고 있습니다. 다만, Post-Training Quantization (PTQ) 기능만 지원하고 있어 모델 정확도 면에서 아쉬운 부분이 많았습니다. Quantization-Aware Training (QAT)과 같은 추가적인 학습을 하려고 해도 위에서 언급한 역변환의 어려움으로 인해 복잡한 개발 과정을 거쳐야 했습니다.
OwLite는 스퀴즈비츠에서 개발한 컴파일러 기술을 기반으로 TensorRT 엔진을 만들어내는 과정에서 손쉽게 QAT가 가능하도록 했습니다. QAT를 위해서는 데이터셋과 학습 시간, 학습 리소스가 더 필요하기는 하지만 동일한 경량화 수준에서 PTQ보다 우수한 성능을 보일 수 있고, 동일한 성능 수준에서는 더 작고 가벼운 모델을 만들어 낼 수 있습니다.
OwLite를 이용하면 TensorRT와 호환이 잘 되도록 하면서도 PyTorch 모델을 수정할 수 있기 때문에 기존의 학습 코드나 data pre-processing, 모델 성능 평가 등의 코드를 그대로 재사용할 수 있다는 점이 OwLite의 독보적인 기술력이자 강점입니다.
다양한 Compression Scheme
사용자는 각 Layer 가 가지는 특성을 바탕으로 수많은 경량화 방법론 중 하나를 적용할 수 있습니다. 예를 들어, 최적의 PTQ Calibration을 위해 absmax, mse, percentile과 같은 알고리즘을 각 Layer 단위로 선택해서 적용할 수 있으며, 그 외에도 QAT backward 알고리즘, per-channel quantization 여부, unsigned representation, grad-scale과 같은 configuration들이 적용될 수 있습니다.
이를 통해, 사용자는 하나의 모델에 일괄적인 방법론을 적용하는 것을 넘어서 모델의 각 부분에 가장 효과적인 기술을 적용할 수 있고 이는 결과적으로 각 모델에 대해 최적의 경량화 방법을 얻어낼 수 있음을 나타냅니다.
뿐만 아니라, OwLite의 출시와 함께 제공되는 scheme들 외에도 지속적인 연구 및 학계와의 교류를 통해 새로운 scheme을 빠르게 추가할 수 있도록 대응하고 있습니다.
기획
사용자를 따라가기
사용하기 편한 도구의 정의는 무엇일까요? 저희는 사용자가 가진 목표에 대한 이해를 바탕으로 만들어진 도구가 편한 도구라고 생각했고, OwLite가 그런 도구가 되기를 바랬습니다. 사용자가 AI경량화 과정에 대해 어떤 멘탈 모델을 가지고 있는지, 목표를 이루기 위해 어떤 의사결정을 하는지, 각 의사결정간의 선후관계와 중요도는 어떻게 되는지를 파악했습니다.
생각의 흐름을 담은 프로젝트, Baseline, Experiment의 위계
OwLite는 (1)문제를 정의하고 (2)해결하고자 하는 방향을 선택한 다음 (3)구체적인 방법을 찾아나가는 생각의 흐름을 반영하도록 기획되었습니다.
AI 모델로 해결하고자 하는 문제, 예를 들면 이미지 내에서 특정 물체를 찾아내는 것이나 텍스트를 생성하는 등의 Task 단위로 업무가 구분될 것이라고 생각하여 해당 위계를 가장 상위의 Project로 구분했습니다.
한 편, Image Classification Task의 ResNet, MobileNet처럼 특정 목표를 위한 여러가지 모델을 테스트 해볼 수 있기 때문에 하나의 Project 내에서도 여러개의 Baseline을 등록할 수 있도록 기획했습니다. Baseline은 경량화되지 않은 모델 자체를 의미하며, 이를 분석하는 것에서부터 경량화 과정이 시작됩니다.
Experiment는 각 Baseline을 기반으로 만들어진 여러 경량화 실험입니다. 사용자는 다양한 경량화 기법을 적용한 Experiment 생성을 통해 목표 수치에 부합하는 경량화 결과물을 얻어낼 수 있습니다.
경량화 과정에서 몰입을 부여하는 GUI 에디터
AI모델을 경량화 하는 과정은 높은 몰입도를 필요로 합니다. 경량화를 위해 많은 보조 지표나 도구들을 사용해야 하며 경우에 따라서는 경량화 할 모델의 구조가 복잡하고 거대해서 시각화에 충분한 공간을 필요로 합니다.
모델 경량화 과정의 몰입도를 부여하고, 복잡한 경량화 과정에서 효율적으로 화면을 사용하고자 OwLite팀은 시스템-에디터의 2-Depth GUI를 기획했습니다. 에디터 화면은 사용자가 온전히 모델의 구조에 집중할 수 있도록 별도의 UI를 가지는데, 모델의 구조와 상세한 인사이트 그리고 경량화를 위한 기능들이 배치되어 있습니다.
좌측에 고정된 ‘OpType Setting’을 통해, 사용자는 해당 모델의 Operator들이 어떤 설정으로 적용되어 있는지를 상시 확인할 수 있습니다. 이와 대조적으로, ‘Layer Setting’은 패널 형태가 아닌, 캔버스 위에서 사용자가 Layer를 클릭했을 때만 나타나는 형식으로 구현되었습니다. 이러한 접근은 임시적인 인터페이스를 통해 ‘OpType’과 ‘Layer Setting’ 간에 명확한 멘탈 모델을 형성하도록 돕습니다. 이는 사용자가 시스템의 작동 방식을 더 직관적으로 이해하고, 효과적으로 경량화 옵션을 적용할 수 있도록 지원합니다.
작업의 흐름을 자연스럽게 따르면서 필요한 기능을 담은 UI
OwLite팀은 경량화 과정을 “가장 쉬우면서도 강력한 방법”에서 시작하여 “복잡하고 세밀한 방법으로 마지막 조율”을 하는 것으로 정의했습니다. 우리에게 시간은 한정된 재화이기에, 시간 대비 효율이 가장 좋은 것에서 시작하여, 완벽을 위해 fine tuning을 하는 순서로 작업이 이루어 질 수 있도록 기획했습니다.
경량화 과정에는 Golden Rule이 없습니다. 스퀴즈비츠의 전문가라고 하더라도 새로운 task와 모델에 대해서는 여러 경량화 방법을 시도해보며 최적의 방법론을 찾아야 합니다. 그렇기 때문에 OwLite는 굉장히 높은 수준의 자유도를 지원하고 있습니다.
하지만, 높은 자유도는 반대로 말하면 경량화가 처음이신 사용자에게는 ‘어디서 시작해야할 지 모르겠다’는 문제를 야기합니다. 그래서 OwLite는 가장 쉬우면서도 강력한 방법으로 Recommended 세팅을 제공합니다. 스퀴즈비츠의 엔지니어들이 다양한 모델 경량화 업무를 통해 얻은 노하우를, 사용자들이 간단히 적용할 수 있도록 버튼으로 만들어 두었습니다. 많은 모델에 범용적으로 적용될 수 있도록 기획된 Recommended 세팅은 경량화의 첫 걸음이 될 수 있을 것입니다.
명확한 가시성, 쉬워지는 경량화
OwLite가 가진 기능의 가시성을 높이기 위한 시도
OwLite는 강력한 경량화를 위해 다양한 경량화 방법론을 지원합니다. Operand (Weight, Input)에 따라 precision, PTQ calibration, QAT backward, per-channel, unsigned, grad scale 의 옵션을 변경 가능합니다. 이를 위해 Table형태의 옵션 편집창이 만들어졌으며, 방법론의 성격에 따라 드롭다운 박스나 토글 버튼으로 편집할 수 있도록 구성했습니다. 숙련도에 따라 가장 기본적 Scheme인 precision만을 조회하도록 UI를 접거나, 상세 설정을 위해 UI를 펼 수 있는 기능을 추가했습니다.
Layer 단위의 경량화 방법론을 설정할 때는 그래프에서 표시된 Input Tensor (A, B) 가 패널의 어떤 Operand (input (A), input (B))와 대응 되는지 표시했습니다. 전후에 배치된 Layer 와 모델 구조를 바탕으로 각각의 Input Tensor의 위치를 파악하면 더 정확하고 세밀한 설정이 가능합니다.
OwLite가 제공하는 정보의 가시성을 높이는 시도
에디터 화면은 세가지 정보를 가지고 있습니다.
(1)시각화된 모델 구조 (2)모델에 대한 부수적인 정보 (3)모델을 경량화하는 방법론에 대한 정보 입니다.
모델 구조를 가시성 있게 시각화 하기 위해 Operator들을 몇 개의 카테고리로 나눈 다음 고유의 색상을 부여했습니다. 각 Layer의 시각화에서는 Operator가 어떤 종류인지, 그리고 Weight, Input 등 Operand 의 정보를 가장 먼저 보여줄 수 있도록 했습니다. 경량화 업무에서는 Operand의 종류와 Tensor Size를 인지하는 것이 매우 중요하기 때문입니다.
OwLite는 PyTorch-ONNX-TensorRT로 이어지는 호환성을 통해 End-to-End Latency 뿐만 아니라 TensorRT에서 측정된 fused kernel 단위의 Latency를 ONNX 구조에 매핑하여 표시하는 기능을 가지고 있습니다. 왼쪽 아래의 Breakdown 기능을 통해 사용자는 어떤 단위의 그룹이 bottleneck이 되고 있는지 알 수 있으며, 해당 영역을 클릭하면 Layer가 강조되는 기능을 통해 그 위치 또한 알 수 있습니다.
재미와 따듯함
기술과 기획의 중요성을 언급했지만, 이와 동시에 서비스는 제공자와 사용자가 소통하는 창구이기도 합니다. 사람과 사람이 맞닿는 부분에서 OwLite는 재미와 따듯함을 담고 있습니다.
깊은 기술을 다루는 도구인 OwLite는 필연적으로 사용 과정에서 사용자에게 여러 메세지를 표시해야 합니다. OwLite에서 사용자가 마주하는 메세지들은 정보, 경고, 에러로 구성되어 있습니다.
정보성 메세지는 OwLite 서비스나 함수의 각 부분이 어떤 동작을 하는지, 그래서 현재 서비스가 어떤 상태이고 다음 상태는 무엇인지 알려줍니다. 무언가를 로딩하거나 기다리는 상황에서도 멈춰져 있지 않다는 정보를 제공하기 위해 CLI 상에서 동적인 요소를 추가했습니다.
또한, 경고와 에러에 해당하는 메세지는 마주해도 당황하지 않도록 (1)원인과 (2)결과 (3)해결책을 담을 수 있게 UX writing을 고려했습니다.
이와 함께, 귀여운 부엉이 로고를 활용하여 로딩 화면에서 부엉이가 눈을 깜빡이는 애니메이션을 추가하고 시스템상의 여러 소식을 전하는 알림센터 역시 부엉이를 이용해 버튼을 만들었습니다. CLI상에서는 부엉이 이모티콘이 대기 상황에 따라 이동하기도 합니다.
AI 경량화라는 멋진 일을 재미있게 할 수 있도록 앞으로도 OwLite팀은 곳곳에 유머와 따듯함을 녹여낼 것입니다.
마치며
OwLite를 개발하며 했던 고민과 결정들을 글로 풀어놓고 나니, 이제야 첫 걸음을 뗀 실감이 납니다. 스퀴즈비츠의 구성원들 모두 함께 좋은 도구를 만들기 위해 적극적으로 참여했습니다. 부족한 부분은 서로 채우고, 강점인 부분은 더욱 발현할 수 있도록 했습니다.
많은 관심과 함께, 더 발전해 나가는 OwLite를 만들 수 있도록 값진 조언 부탁 드립니다.
