[AI 경량화] tfmot 패키지에 대해 알아보자 (tensorflow-model-optimization)

---

서  론

tfmot 패키지는 TF Team에서 공식으로 배포하고 있는 AI 모델 최적화 패키지이다.

아직 버전이 0.3.0 밖에 안 되는 beta 버전이지만, 이거라도 없으면 직접 AI 경량화를 진행해야 하는 만큼, 상당히 유용하다.

 

사실 AI경량화 혹은 최적화 부문은 시장도 너무 작고,

원하는 기업도 별로 없다.

 

국내에선 대표적으로 매스프레소의 콴다 앱, AI 경량화 전문 노타 등을 제외하고는,

대기업에서조차 AI 경량화 인재에 국한하여 채용을 진행하지는 않는다. (2020.06.20 기준)

 

사실 생각해보면 CVPR2020 1,467건의 Accepted Paper 중 고작 23건만이 AI Optimization에 관한 것이었으니… 당연할지도 모른다.

하지만 작년 CVPR2019에서 10건밖에 AI Optimization 관련 Paper가 없었던 것을 생각하면, 2.3배라는 놀라운 수치로 관심도가 올라간 것이라고 생각할 수도 있다.

 

실제로 요즘 AI Image Classification / Object Detection 관련 세미나에 가보면,

학계에서는 매우매우 정확도가 높은, 그러나 동시에 너무나도 복잡하고 무거운 모델을 자랑하고,

현업 종사자들은 Q&A 시간에 해당 모델이 너무 무거운데, 디바이스 온보딩 AI로 넣을 때의 방안은 준비되어 있냐는 질문을 하는 모습을 접할 수 있다.

 

음…

물론 내가 들어봤던 대답은 항상 '없다'였다.

 

사실 모델링을 연구하는 학자들이 굳이 모델 경량화까지 책임져야할 이유가 없기도 하고.

 

그래서 결국 상품화를 위한 모델 경량화 및 AI 최적화는 엔지니어의 몫인 것이다.

날이 갈수록 점점 무거워지고 복잡해지는 모델에 대응하여 그만큼 AI 경량화 기술은 점점 더 필요해질 것이고,

결국엔 AI 경량화 팀이 따로 생길 지경까지 가지 않을까?

 

……라는 기대를 품어본다.

아마 그때가 되면 이 블로그의 AI경량화 포스팅도 더 많은 활용을 받겠지?

---

tfmot의 구조

tfmot 패키지의 현재 (v0.3.0) 내부 구조는 아래와 같다.

tfmot
└─┬── quantization
  │    └─ keras
  │        ├─ QuantizeConfig
  │        │    ├─ get_activations_and_quantizers(layer)
  │        │    ├─ get_config()
  │        │    ├─ get_output_quantizaers(layer)
  │        │    ├─ get_weights_and_quantizaers(layer)
  │        │    ├─ set_quantize_activations(layer, quantize_activations)
  │        │    └─ set_quantize_weights(layer, quantize_weights)
  │        ├─ quantize_annotate_layer(to_annotate, quantize_config=None)
  │        ├─ quantize_annotate_model(to_annotate)
  │        ├─ quantize_apply(model)
  │        ├─ quantize_model(to_quantize)
  │        ├─ quantize_scope(*args)
  │        └─ quantizers
  │             ├─ LastValueQuantizer(num_bits, per_axis, symmetric, narrow_range)
  │             ├─ MovingAverageQuantizer(num_bits, per_axis, symmetric, narrow_range)
  │             └─ Quantizer(Object)
  └── sparsity
       └─ keras
           ├─ ConstantSparsity(target_sparsity, begin_step, end_step=-1, frequency=100)
           ├─ PolynomialDecay(initial_sparsity, final_sparsity, begin_step, end_step, power=3, frequency=100)
           ├─ PrunableLayer()
           │    └─ @abc.abstractmethod
           │       get_prunable_weights()
           ├─ prune_low_magnitude(to_prune, pruning_schedule=pruning_sched.ConstantSparsity(0.5, 0), block_size=(1, 1), block_pooling_type='AVG', **kwargs)
           ├─ prune_scope()
           ├─ PruningSchedule(step)
           ├─ PruningSummaries(log_dir, update_freq='epoch', **kwargs)
           ├─ strip_pruning(model)
           └─ UpdatePruningStep()

TF Team에서 만들었지만, 역시 버전이 너무 낮아서 그런지 내부 구조와 내포되어 있는 함수의 양은 매우 적다.

 

이 함수들을 크게 분류하자면,

1. AI 모델 Pruning(프루닝) - tfmot.sparsity로 표현된다.

2. AI 모델 Quantization(양자화) - tfmot.quantization으로 표현된다.

두 가지로 나눌 수 있다.

 

각각의 경량화법은 모델의 사이즈를 축소시키고, inference latency를 향상시킨다.

대신, 정확도는 희생되며 loss가 늘어난다.

 

이 중 Pruning은 Training Callback Pruning과 Post-training Pruning을 지원하고,

Quantization은 Quantization Aware Training만을 제공한다.

 

즉, Pruning은 학습 전과 후를 모두 지원하고, Quantization은 학습 전만 지원한다는 뜻이다.

그렇다면 학습 후에는 Quantization을 못하는가? 하면, 그것은 또 아니다.

Post-training Quantization은 tfmot 패키지가 아닌 TFLite 패키지에서 제공하니, 추후 포스팅에서 다루도록 하겠다.

---

Pruning

 

Pruning은 원래 가지치기라는 뜻을 지니고,

실제로 구글에서 Pruning을 검색하면 아래와 같은 이미지가 더 많이 검색된다.

tfmot 패키지에서 Pruning이 하는 역할 또한 똑같다.

Heavy-weighted 모델에서 Pruning을 통해 weights를 Prune(가지치기)하고, 그 원리로 모델을 경량화 시킨다.

 

학습 전 Pruning은,

우선 1차로 Compile이 끝난 모델을, fit을 시키며 학습시킬 때 중간중간 Callback 함수를 통해 매 레이어마다 Pruning 시킨다.

 

학습 후 Pruning은,

학습이 완료된 모델을 통째로 Pruning 시킨다.

 

이 중 학습 전 Pruning은 너무나도 불편하고, 프로그래밍 하기도 불편하고, 모델에 대한 숙지 없이 Pruning을 진행한다면 너무 큰 정확도 손실을 야기할 수 있다.

Tensorflow Official Doc에서는 attention mechanism 등의 핵심 레이어는 Pruning하지 말라고 하고,

이후 레이어(Later layer)를 Pruning하는 것이 초기 레이어를 Pruning하는 것보다 좋다고 권장한다.

 

학습 후 Pruning은 프로그래밍도 간단하고, 모델에 대한 숙지도 필요없지만,

향상되는 inference latency 폭과 축소되는 모델 사이즈의 폭이 작다.

---

Quantization

 

원래 Quantization(양자화)란, 본디 아날로그 통신을 디지털 통신으로 변환할 때 사용되는 단어였다.

본디 Digital Signal Processing쪽에서 쓰이던 용어였고,

그 뜻은 아날로그 신호의 표본화된 진폭치로부터 특정 이산치(대표치) 만으로 바꾸는 과정이다.

 

이는 AI경량화에서도 비슷하면서도 다른 뜻으로 통용되는데,

대표적으로 tfmot이 제공하는 Quantization은 Float32, Float16, UINT8 등의 양자화가 있다.

학습된, 혹은 앞으로 학습할 모델의 가중치를 해당 자료형으로 양자화시키는 것이다.

 

물론 Float32가 가장 정확도가 높지만 무겁고,

UINT8은 정확도의 손실이 조금 있지만 훨씬 가볍다.

 

최신 sota모델인 EfficientNet B0 기준, tfmot이 아닌 TFLite로 Post-training UINT8 Quantization을 진행했을 때,

추론 속도가 0.09 FPS에서 1.20 FPS까지 올랐으니 아무 것도 안 했을 때보다 inference latency가 약 13.3배 가량 올라가게 되었다.

 

Post-training Quantization보다 Quantization Aware Training이 훨씬 프로그래밍하기 복잡하고,

적용할 수 있는 모델도 한정적이다 보니,

Quantization Aware Training의 결과가 훨씬 좋은 편이다.

---

결   론

tfmot은 아직 베타 버전인 AI 경량화 패키지이며,

제공하는 기능으로는 학습 전 Pruning / 학습 후 Pruning / Quantization Aware Training이 존재한다.

 

위 기능들은 정확도를 조금 희생하는 대신 Inference Latency와 모델 사이즈를 혁신적으로 향상시키며,

여러 가지 경량화법을 중복으로 적용하거나, 단일로 적용할 수도 있다.

 

오늘 포스팅에서는 tfmot에 포함된 함수와 그 구조, 그리고 각각의 아주 간단한 이론적 배경만 알아보았다.

다음 포스팅에서 TFLite를 소개하는 것까지 끝난 후에는,

Pruning, Quantization의 더 자세한 이론과, 이를 코드로 구현하는 법을 공유하고자 한다.

 

다음 포스팅도 기대된다.

---