[Paper Review] MobileOne: An Improved One millisecond Mobile Backbone

 

Paper : https://arxiv.org/pdf/2206.04040.pdf

GitHub : https://github.com/apple/ml-mobileone

GitHub - apple/ml-mobileone: This repository contains the official implementation of the research paper, "An Improved One millis

 

 

오늘 읽어볼 논문은 Apple에서 공개한 MobileOne이며, CVPR 2023에 소개될 내용이라고 하네요, 실제 논문이 공개된 시점은 2022년 입니다. 

 

논문에 따르면 모바일 환경을 위한 backbone 구조는 종종 FLOP 또는 파라미터 수와 같은 metric에 맞춰서 나오는데, 이러한 방식은 실제로 모바일 기기에 배포될 때 network latency가 잘 고려가 안된다고 하네요. 그래서 애플에서는 mobile friendly network 구조인 MobileOne backbone을 제안합니다. 이를 위해 ImageNet에서 75.9%의 정확도로 iPhone 환경에서 1ms 미만의 추론 시간을 달성한다고 하네요. 본 논문에서 제안한 모델 중 best model은 MobileFormer 방법과 비교했을 때 유사한 성능을 보이면서 약 38배 정도 빠르다고 하네요. 그리고 비슷한 latency에서 EfficientNet 보다 2.3% 정도 더 정확하다고 합니다. 이와 관련하여 코드도 물론 공개가 되어 있습니다. 

 

Introduction

본 논문에 따르면 skip-connection 이나 branch와 같은 disconnect 즉, 매개 변수가 없는 parameter-less 작업들은 상당한 memory access cost를 갖는다고 합니다. 그래서 본 논문에서는 on-device latency를 고려하여 아키텍처를 설계하고, optimization bottleneck 을 통해 latency를 개선하는 것을 목표로 합니다. optimization bottleneck 현상을 해결하기 위해 train time 및 inference time에 따라 아키텍처를 분리합니다. 즉 train 할 때에는 선형적으로 과도하게 매개변수화 된 모델(linearly over parameterized model)을 사용하고, inference 할 때에는 linear 구조(branch나 skip-connection이 없는)를 다시 매개변수화 하여 사용합니다. 이미 작은 모델이 과도하게 정규화 되는 것을 방지하기 위해 train 과정에서 정규화 과정을 동적으로 완화 시켜 optimization bottleneck 현상을 완화시킵니다. 이러한 구조는 memory accecss cost가 낮기 때문에 신경망에 더 넓은 레이어를 통합하여 더 많은 표현을 나타낼 수 있도록 할 수 있습니다. 

 

예를 들면 MobileOne-S1 구조에서는 4.8M개의 파라미터, 0.89ms의 latency가 발생하는 반면, MobileNet-V2 구조에서는 3.4M개의 파라미터를 가지지만(29.2% 적음) 0.98ms 만큼의 시간이 걸립니다. 즉, 파라미터 수가 적더라도 시간이 더 오래 걸리는 결과를 확인할 수 있습니다. 정확도도 MobileOne이 3.9% 정도 낫다고 합니다. 

 

MobileOne의 Main Contribution은 아래와 같습니다. 

 

• 모바일 장치에서 1ms 이내 실행되고, Image Classification에서 SOTA를 달성함, 데스크탑 CPU 및 GPU에서도 일반화 됨 
• 모바일 환경에서 최근 등장한 신경망들이 왜 높은 latency를 가지는지 분석 
• re-parameterizable branch의 효과와 정규화의 동적 완화를 분석

 

 

 

 

Related Work

 

파라미터 수 등 일반적으로 최적화 

• SqueezeNet
• Mobile-ViT

 

부동 소수점 연산(FLOP)의 수를 최적화 

• MobileNeXt
• ShuffleNet-V1
• GhostNet
• MixNet

 

FLOP 최적화와 동시에 depth, width, resolution를 복합 스케일링

• EfficientNet
• TinyNet

 

latency를 위해 직접 최적화 

• MNASNet
• MobileNet-V3
• ShuffleNet-V2

 

Dehghani et al. [9] 방법은 FLOP 및 파라미터 수가 latency와 상관 없음을 연구하였으며, MobileFormer 방법은 FLOP을 최적화 하여 낮은 FLOP에서 효율적인 CNN을 능가하였습니다. 또한 MobileNet-V3에서는 특정 플랫폼에 대해 최적화된 activation function인 Hard-Swish를 도입합니다. 하지만 이러한 기능은 다른 플랫폼으로 확장하는데 어려울 수 있다고 합니다. 따라서 본 논문에서는 다른 플랫폼에서 모두 사용할 수 있도록 기본 연산자들을 사용합니다. 

 

 

 

Method

 

실제로 아래 표 1과 같이 상관관계를 계산해보면, 모바일 환경에서는 FLOP의 수는 중간 정도 상관 관계가 있고, Parameter의 수는 그닥 크게 상관 있지 않다는 것을 보여주며, CPU 환경에서는 상관관계가 더 낮다는 것을 보여줍니다. 

 

Key Bottlenecks

Activation Function 

활성화 함수가 latency에 미치는 영향을 분석하기 위해 본 논문에서는 30 layer conv로 이루어진 신경망을 구성하고 벤치마킹 작업을 하였습니다. 표 2를 보면, SE-ReLU 및 Dynamic Shift-Max, DynamicReLU와 같은 최근 활성화 함수들은 synchronization cost가 많이 든다고 합니다. DynamicReLU 및 Dynamic Shift Max는 MicroNet과 같은 매우 낮은 FLOP 모델에서 상당한 정확도 향상을 보여주었지만 시간이 상당히 들 수 있다고 합니다. 그래서 MobileOne 에서는 ReLU 만 사용합니다. 

 

 

 

Architectural Blocks

runtime 성능에 영향을 미치는 두 가지 핵심 요소는 memory access cost와 병렬처리를 얼마나 하는가 입니다. multi branch 구조 같은 경우 graph의 다음 tensor를 계산하기 위해 각 branch의 activation을 저장해야 하기 때문에 memory access cost가 크게 증가한다고 합니다. 이러한 메모리 병목 현상은 신경망에서 분기를 최소화 하여 피할 수 있습니다. Squeeze-Excite block에서 사용되는 global pooling 과 같은 작업도 synchronization cost으로 인해 실행 시간에 영향을 미치게 됩니다.

그래서 본 논문에서는 skip connection 및 squeeze-excite block을 없앱니다. 이에 따른 표 3에 나타나 있습니다. 특히 이 SE block은 정확도를 향상시키기 위해 MobileOne의 가장 큰 버전에 들어간다고 합니다. 

 

 

 

MobileOne Architecture

MobileOne Block

MobileOne Block은 depthwise 및 pointwise layer로 나눠지는 conv layer로 설계되었고, 정확도 향상을 위해 over parameterization branch를 도입합니다. basic block은 MobileNet-V1 3x3 depthwise conv 를 기반으로 합니다. 그 다음 그림 3과 같이 구조를 복제하는 branch와 함께 batchnorm을 사용하여 re-parameterizable skip connection을 도입합니다. inference 과정에서는 branch가 없습니다.

 

 

 

 

re-parameterizable branch의 유무에 따른 성능은 아래와 같습니다.

 

 

 파라미터 계수 $k$는 1에서 5까지의 hyperparameter이며, 이에 따른 실험 결과는 아래와 같습니다. 

 

Model Scaling

본 논문에서는 5가지 width scale을 사용합니다. 또한 모바일에서는 입력 해상도를 크게 하는 것이 효과적이지 않기 때문에 그에 대한 실험은 하지 않습니다. 

 

 

Training

large model과 달리 small model은 over-fitting을 방지하기 위해 정규화가 "덜" 필요합니다. 그래서 학습 초기에 weight decay를 갖는 것이 중요합니다. 또한 weight decay regularization을 제거하는 것 보다 학습 과정에서 loss를 annealing 하는 것이 효과적이라는 사실을 발견했다고 합니다. 그래서 learning rate에 cosine schedule을 사용합니다. 또한 동일한 schedule을 사용하여 weight decay coefficient를 annealing 하고, progressive learning curriculum을 사용합니다. 그에 따른 실험 결과는 다음과 같습니다. 

 

 

Benchmarking

모바일 장치에서 정확하게 실행 시간을 측정하는 것은 어려운 작업이며, 본 논문에서는 iPhone 12를 기준으로 swift를 사용하여 iOS 애플리케이션을 개발하여 시간을 측정 했다고 합니다. 애플리케이션에서 모델은 CoreML을 이용했다고 하네요. 시간 측정은 1000번 정도 반복해서 측정했고, 다른 모든 앱은 다 종료시켰다고 하네요. CPU 환경에서 실행할 때는 GHz - Intel  Xeon Gold 5118, RTX-2080Ti의 Ubuntu 환경에서 TensorRT 라이브러리를 사용하여 모델을 100회 정도 반복하여 실행했다고 합니다. 

 

 

 

Experiments

Image Classification on ImageNet-1K

본 논문에서는 128 만 개의 학습 이미지와 1,000개 클래스의 50,000개 이미지로 구성된 validation 세트로 구성된 ImageNet 데이터 세트에서 MobileOne 모델을 평가합니다. 모든 모델은 8개의 NVIDIA GPU가 있는 시스템에서 PyTorch 라이브러리를 사용하여 스크래치 학습됩니다. 모든 모델은 Momentum optimizer와 함께 SGD를 사용하여 배치 크기 256으로 300 epoch 동안 학습됩니다.

 

모든 모델에 대해 smoothing factor를 0.1로 설정한 cross entropy loss과 함께 label smoothing regularization을 사용합니다.

 

initial  learning  rate은 0.1이고 cosine schedule을 사용하여 annealing됩니다. Initial weight decay coefficient는 $10^{-4}$로 설정되고 논문[42]에 설명된 것과 동일한 cosine schedule을 사용하여 $10^{-5}$로 annealing됩니다.

 

또한 AutoAugment를 사용하여 MobileOne S2, S3, S4를 학습합니다. auto augmentation의 강도와 이미지 해상도는 논문[56]에 소개된 것처럼 학습 중에 점진적으로 증가합니다.

 

MobileOne S0, S1 모델의 경우, standard augmentation(random resized cropping and horizontal flipping)을 사용합니다. 

 

또한 모든 버전을 학습하기 위해 decay constant = 0.9995 값을 사용하는 EMA(Exponential Moving Average) weight averaging 방법을 사용합니다. 

 

모든 테스트는 224x224 해상도에서 진행되며, Object Detection, Image Segmentation 분야에서도 효과가 있음을 입증합니다. 

 

 

마이크로 버전 성능도 있습니다. 

 

CPU, GPU, TPU 환경에서의 성능은 아래와 같습니다. 

 

 

이외에 부록에서는 MobileNet V3에서 H-swish 연산자가 특정 하드웨어 플랫폼에 최적화 되어있고 다른 플랫폼에서는 최적화 되어있지 않다고 언급하고 있습니다. Howardet al. [30] 논문에서는 H-swish가 플랫폼 별 최적화가 적용될 때 ReLU와 유사한 성능을 얻을 수 있다는 것을 입증했기에 H-swish layer를 ReLU layer로 대체하여 시간을 측정했다고 하네요. 부록에 볼만한 내용이 많은 것 같습니다.