TFLite, 너 T야? F야?. 안녕하세요. AI 추론 최적화 엔진을 만들고 있는 Optima 팀의… | by ENERZAi | Jan, 2024

안녕하세요. AI 추론 최적화 엔진을 만들고 있는 Optima 팀의 모세웅(Sewoong Moh)입니다. 요즘 관심있는 사람에게 MBTI를 물어보듯이 이 글에서는 AI 추론 최적화에서 많이 사용되는 엔진인 TensorFlow Lite(TFLite)에 대해서 보다 더 심도 있게 알아보려고 합니다.

TFLite은 모바일, 마이크로컨트롤러 및 기타 엣지 기기에 모델을 배포하여 On-device AI를 실현할 수 있도록 도와주는 엔진입니다. On-device AI를 중점적으로 지원하기 때문에 짧은 Latency로 구동되도록 작은 바이너리 크기로 만들어 주는 것에 집중하고 있는 엔진입니다. 원래 해당 엔진은 Google Brain에서 내부 활용 목적으로 개발하였으나 2015년부터 오픈 소스로 제공되고 있습니다. 현재는 40억 개 이상의 장치에서 실행될 정도로 범용적인 추론 엔진 중 하나입니다. 이처럼 많은 곳에서 사용되는 TFLite는 다음과 같은 특징을 가지고 있습니다

지원 내용

• 멀티 플랫폼: Android, iOS, Embedded Linux, Microcontroller를 지원합니다.
• 언어: Java, Swift, Objective-C, C++, Python 와 같은 여러 언어를 지원합니다.
• 높은 성능 : 하드웨어 가속과 모델 최적화 지원합니다.

특징

• Latency: 서버까지의 왕복이 없기 때문에 지연 시간이 적습니다.
• Privacy: 네트워크를 사용하지 않아 개인 정보가 외부로 유출되지 않습니다.
• Connectivity: 서버로 데이터를 보내지 않기 때문에 인터넷에 연결이 필요하지 않습니다.
• Size: 모델 및 바이너리 크기가 축소됩니다.
• Power Consumption: 효율적인 추론을 지원하며, 네트워크 연결이 불필요하기 때문에 전력 소비가 적습니다.

TFLite와 FlatBuffer

TFLite는 .tflite 파일 확장자로 표시되며 ‘FlatBuffer’라고 불리는 특수하고 효율적인 휴대용 형식을 사용합니다. Flatbuffer는 flatc라는 컴파일러를 이용하여 직렬화 또는 역직렬화가 가능합니다. 마찬가지로 .tflite 파일 또한 역직렬화가 가능합니다. 다시 말해 구문 분석이나 압축 풀기 없이 직렬화된 데이터에 액세스가 가능합니다. 또한 Flatbuffer는 아래와 같은 장점을 추가적으로 가지고 있습니다.

• 사람이 읽을 수 있는 모델 설명과 기계가 읽을 수 있는 데이터가 포함된 메타데이터를 추가합니다.
• TF의 Protocol Buffer보다 크기와 Latency가 작아 최소한의 컴퓨팅 및 리소스를 사용합니다.
• Protocol Buffer에 비해 Small Code Footprint를 가지고 있기 때문에 크기가 작습니다.
• 별도의 Parsing, Unpacking 과정 없이 데이터에 직접적으로 접근하기 때문에 빠른 추론 속도를 보입니다.

TFLite와 Quantization

학습된 TF 모델의 추론을 지원하는 Quantization은 아래와 같이 3개 종류가 있습니다. Quantization를 사용하게 되면 Trade-off(ex. Accuracy가 낮아지면 전력 효율성이 높아짐)가 발생하게 됩니다. 따라서 사용자가 해결하고 싶은 우선순위에 대한 구체적인 기준이 있어야 합니다. TF팀은 사용자가 가장 적합한 Quantization 방법을 선택할 수 있도록 아래와 같은 결정 트리를 제공합니다.

Dynamic Range Quantization

• 대표(Representative) 데이터 세트를 제공할 필요 없이 메모리 사용량을 줄이고 계산 속도를 높이기 때문에 TF에서 권장합니다.
• CPU에서 작동하며 모델을 4배 더 축소하고, 2~3배 빠르게 합니다.

Full Integer Quantization

• 모든 모델의 Latency가 개선되며, 최대 메모리 사용량 감소됩니다. 또한, 특정 하드웨어 장치 및 Accelerator와 호환성이 높아집니다.
• CPU / MCU에서 작동하며 모델을 4배 축소하고, 3배 빠르게 합니다.

FLOAT16 Quantization

• 상대적으로 다른 Quantization보다 압축률은 작지만 성능 손실이 거의 없습니다.
• CPU / GPU에서 작동하며 가속을 지원합니다.

TFLite Model을 사용하는 법

TFLite Model Maker

• 이 경우 사용자의 데이터를 통해 모델 훈련 할 수 있습니다. 지원되는 작업은 이미지 분류, 객체 감지, 텍스트 분류, BERT 질문 답변, 오디오 분류, 추천 등이 있습니다. 해당 기능은 전이 학습을 활용하여 필요한 훈련 데이터 양을 줄이고 전체 훈련 시간을 줄일 수 있는 장점이 있습니다.

TFLite Converter

• TF Model을 TFLite Model로 Converting 해야하는 경우에 사용됩니다. Accuracy (혹은 설정한 별도의 Metric)의 손실이 없거나 최소한의 손실로 Model Size를 줄일 수 있습니다. 또한 Latency를 줄이기 위한 Optimization, Quantization 옵션을 적용 할 수 있습니다.
• 최적화와 Quantization를 적용하여 모델 크기와 Latency를 줄입니다. TFLite는 TF 모델을 사용하여 .tflite 파일 확장자로 식별되는 최적화된 FlatBuffer 형식을 생성합니다. 또한 해당 Converter에는 모델을 변환하기 위한 Python API가 포함되어 있습니다.

TF Hub

• “TFLite Model을 사용하는 법” 중에 가장 간단한 방법으로는 이미 잘 만들어진 모델을 TF Hub에서 무료로 다운 받아 사용하는 방법이 있습니다. 사용자는 다운 받은 모델을 그대로 사용하거나 전이 학습을 통하여 사용자의 목적에 맞게 자유롭게 수정 할 수 있습니다.
• 2023년도 전까지는 자체적으로 TF Hub 홈페이지가 있어, TF/TFLite 관련된 모델만을 다운받을 수 있었습니다. 2023년 말부터는 Kaggle 홈페이지로 흡수되어 운영되고 있습니다. 해당 홈페이지에서 사용자의 편의성을 높이기 위해서 더 다양한 JAX, ONNX, MXNet 등 다양한 프레임워크를 지원하는 모델을 무료로 다운 받을 수 있습니다. 구체적으로 307개의 Backbone을 통하여 만들어진 2,186개의 AI 모델을 다양하게 지원합니다.

TFLite와 TF는 무슨 관계?

위 내용들에서도 말씀드렸다시피 TF는 ML 모델을 구축하고 훈련하는 데 사용되지만 TFLite는 이미 훈련된 모델에 대해 예측을 수행하는 기능을 제공하도록 되어 있습니다.

• TFLite는 위에서 말씀드린 Quantization을 바탕으로 훈련된 모델을 최적화하여 필요한 메모리 사용량과 신경망 활용에 따른 계산 비용을 줄입니다. TF도 TFLite와 동일하게 Quantization를 제공하지만 Fake Quantization이기 때문에 오히려 Overhead가 증가하는 한계가 있습니다.
• 모델뿐만 아니라 TFLite는 자체적으로 TF보다 가볍습니다. 이는 TFLite는 모델 추론에 집중되어 있기 때문에 최소한의 Runtime을 가지고 있으며 비교적 설치하는 라이브러리 용량도 작기 때문입니다.

학습된 모델은 학습에 활용된 여러 연산자(Operator)들로 구성이 되어 있습니다. 이러한 연산자들은 모델 추론을 위해 변환의 과정을 반드시 거쳐야하는데, TFLite에서는 아래 3가지 종류의 모델에 대해서 연산자 변환을 지원합니다.

• TFLite 내장 연산자로만 만들어진 모델
• TFLite 내장 연산자와 및 TF 코어 연산자로 만들어진 모델
• TFLite 내장 연산자가 있는 모델, TF 코어 연산자 또는 Custom Operator 사용한 모델

TF 공식 홈페이지에서도 가능하면 TFLite 내장 연산자 사용을 지향하고 Custom Operator를 포함하는 사용은 지양하도록 안내합니다. 왜냐하면 Custom Operator가 변환되면, 내장 TFLite 라이브러리와 외부간 불일치 되어 전체적인 모델의 크기를 증가시키기 때문입니다. 즉, 제한된 리소스를 가지고 있는 엣지 환경에서 사용되는 경우에 치명적인 문제가 될 수 있습니다.

XNNPACK은 TFLite 모델의 추론 속도를 높이기 위한 엔진입니다. XNNPACK은 ARM, WebAssembly, X86 플랫폼에서 부동 소수점 작업에 대한 구현에 집중하여 CPU에서의 빠른 추론 성능을 지원합니다. 평균적으로 2.3배 더 빠른 추론을 가능케하며, FP32 모델과 FP16, Quantization를 지원합니다.

XNNPACK의 장점

• ARM NEON, 어셈블리 최적화: 모바일 장치에서 최고의 성능을 제공하기 위해 모든 연산자는 ARM NEON에 최적화되었으며, 가장 중요한 연산자는 휴대폰에서 일반적으로 사용되는 ARM 코어에 맞게 어셈블리 수준으로 구현되어 있습니다. 구체적으로 Pixel 2의 Cortex-A53/A73 및 Pixel 3의 Cortex-A55/A75에 맞게 최적화가 되었습니다.
• Operation Fusion: TensorFlow Lite 연산자를 하나씩 실행하는 대신 전체 계산 그래프를 살펴보고 Operation Fusion을 통해 최적화합니다. TFlite에서 PAD 연산자와 VALID CONV_2D 연산자 두 개로 분리되어 표현되어 있지만 XNNPACK에서는 자체적으로 Operation Fusion이 가능함을 감지하고 두 연산자를 PAD가 포함된 단일 CONV_2D로 실행합니다.
• 다양한 최적화: 메모리 레이아웃 최적화, Vectorization, 멀티 스레드 추론, 희소 커널 추론 등 다양한 최적화를 제공합니다.

XNNPACK의 단점

• XNNPACK은 부동 소수점 계산 전용으로 설계되었기 때문에 고정 소수점 Quantization 모델을 지원하지 않습니다.
• 동적으로 할당된(kTfLiteDynamic 할당 유형 포함) 입력 및 출력은 지원되지 않습니다.
• XNNPACK팀에서 어셈블리 수준으로 각 Device마다 최적화하였습니다. 하지만 새로운 Device가 출시되는 경우 사용자는 XNNPACK팀이 해당 Device를 지원할 때까지 최적화를 사용할 수 없습니다. 비록 사용자가 직접 Custom Layer를 추가할 수 있지만 난이도가 높기 때문에 원하는 최적화로 사용하기 어렵습니다.
• 같은 연산이더라도 Input 크기나 Attribute 등에 따라서 최적의 구현 방식(ex.구현 알고리즘, Vector 크기, Unrolling Size 등)이 다를 수 있으나 고정이 되어 있어 최적화 효과를 극대화하기 어렵습니다.

결론

XNNPACK은 TensorFlow Lite 모델의 추론 속도를 높이는 데 매우 효과적이며 특히 모바일 장치에서 뛰어난 성능을 제공합니다. 하지만 고정 소수점 Quantization된 모델이나 다양한 연산자를 지원하지 않는다는 점에서 사용자의 편의성이 떨어집니다.

1. TFLite는 주요 연산을 지원함으로써 최적화된 작은 바이너리 사이즈를 구현합니다. 다시 말해, TF에서는 지원하지만 TFLite에서는 지원하지 않는 연산이 있다는 뜻입니다. Type 측면에서 보면 대부분의 TFLite의 연산들은 FLOAT32, UINT8, INT8을 대상으로 하기 때문에 FLOAT16 또는 STRING을 위한 많은 연산들은 아직 지원이 안됩니다. 자세한 예시는 아래와 같습니다.

• 예를 들어 PyTorch의 경우 adaptive pooling이 지원되는 반면, TF에는 global averaging pooling 밖에 없는 제약조건이 존재합니다. 따라서 PyTorch에서 만든 모델을 TF로 만들려면 padding을 추가해야 합니다. 이 경우 Accelerator에서 대개 성능 하락을 야기하는 문제가 발생됩니다.

2. TFLite은 모바일 애플리케이션을 염두에 두고 설계되었기 때문에 Intel 및 AMD x86 프로세서에 최적화되지 않았습니다. 현재 TFLite는 대부분 Neon_2_SSE 브릿지를 통해 ARM Neon 명령을 SSE로 변환하는데 의존합니다.

3. 서로 다른 버전의 연산을 사용하는 것 외에도 부동 소수점 모델과 Quantization된 모델 사이에는 변환 방식에서도 차이가 있습니다.

4. PyTorch와 호환성 문제가 있을 수 있습니다.

• PyTorch 모델을 ONNX로 변환해서 TFLite로 변환하는 방법은 있지만 연산상 변환이 1대1로 되는 경우가 적습니다. 대표적인 예로 PyTorch에서 사용하는 Pixel Shuffle이 있습니다. TF에서는 Depth to Space로 대응이 되는데요 아래 그림과 같이 서로 다른 Layout으로 Depth(Channel)을 Spatial Dimension으로 옮깁니다. 자세한 사항은 아래와 같습니다.