[Paper 리뷰] Once-For-All: Train One Network and Specialize it for Efficient Deployment

Once-For-All: Train One Network and Specialize it for Efficient Deployment

---

• 다양한 edge device에서 자원 제약을 만족하는 효율적인 inference의 어려움
• Once-For-All (OFA)
  • training / search를 분리하여 cost를 줄임
  • OFA에서 추가적인 training 없이 sub-network를 얻을 수 있음
• Progressive Shrinking
  • OFA network training을 위해 제안 -> depth/width/kernel/resolution pruning
• 논문 (ICRL 2020) : Paper Link
  • OFA의 개선버전 : CompOFA 리뷰

---

1. Introduction

• model size의 증가는 hardware가 가진 latency, energy 등의 제약조건으로 인해 다양한 hardware 적용을 어렵게 함
  • For example,
    - 다양한 hardware 리소스를 사용하면 최적 neural network architecture가 달라짐
    - 동일한 hardware를 쓰더라도, 배터리, workload가 달라지면 model도 달라짐
• NAS를 하더라도 모든 시나리오에 맞는 특화된 DNN을 만드는 건 비용이 많이 듦
  • 네트워크 디자인 과정을 계속해서 반복해야 하고, 설계된 네트워크를 처음부터 다시 train 해야 하기 때문

-> 그래서 Once-For-All (OFA)를 제안

 

< Overall of OFA >

• inference 할 때는 OFA의 일부 파트만 선택하여 수행되고, retraining 없이 다양한 depth, width, kernel, resolution을 지원
• training stage / neural architecture search stage를 분리 
  • model training stage
    - OFA에서 파생된 모든 sub-network의 정확도를 향상하는 것을 목표로 함
  • model search (specialization) stage
    - sub-network의 subset을 샘플링해서 accuracy predictor와 latency predictor를 train
    - 주어진 hardware 제약조건에 따라서 predictor가 architecture search
    -> 이렇게 해서 O(n)에서 O(1)까지 떨어트림
• OFA training 할 때 sub-network간 가중치의 joint optimization이 필요한 문제
  • 서로 다른 sub-network가 서로 방해해서 training process를 비효율적으로 만들기 때문

-> 그래서 Progressive Shrinking Algorithm을 OFA에 도입

 

< Overall of Progressive Shirinking >

• OFA network를 직접 optimize 하는 대신, 가장 큰 network를 training
  • 가장 큰 network = (max depth, max kernel, max width) network
• 이후 작은 sub-network가 큰 network와 가중치를 공유하도록 OFA network를 progressively fine-tune
  • better initialization : 큰 sub-network에서 중요한 가중치를 선택할 수 있기 때문
  • distillation : 작은 sub-network 추출해서 training efficiency 향상

-> 하나의 pruned network로 만드는 것 보다 모든 sub-network의 정확도를 유지하는 것이 목표

OFA performance

2. Method

- Problem Formalization

• $W_{o}$ : OFA Network 가중치
• ${ arch_{i} }$ : architectural configuration
• $C(W_{o}, archi_{i})$ : $W_{o}$ 가중치의 OFA network에서 $archi_{i}$ configuration의 sub-network를 select
• Training 목표
  : 각 sub-network가 개별적으로 학습하는 것과 동일한 정확도를 가지도록 $W_{o}$ 최적화

OFA Formalization

- Architecture Space

• CNN의 네가지 차원 : Depth, Width, Kernel, Resolution
• 일반적인 CNN model과 마찬가지로, feature map size는 점점 줄어들면서 channel number가 늘어나는 구조
  - feature map size가 줄어드는 경우 first layer의 stride=2
  - 나머지 모든 layer의 stride=1
• Elastic Depth : 각 unit이 가지는 임의의 layer 개수 / Elastic Width : 임의의 channel 개수
  Elastic Kernel Size : 임의의 kernel size / Elastic Resolution : 임의의 input image size
  • For example,
    - image range : 128~224, stride=4 -> 25개 (128, 132, 136, …, 220, 224)
    - depth : {2,3,4}
    - width expansion ratio : {3,4,6}
    - kernel size : {3,5,7}
    - 5 unit
    - $((3\times3)^2+(3\times3)^3+(3\times3)^4)^5 \approx 2 \times 10^{19}$ -> 이게 25개 씩 (for each resol.)
  • OFA에서는 모든 sub-network는 같은 가중치를 공유하기 때문에 7.7M 밖에 안됨

- Training OFA Network with Progressive Shrinking

• Naive Approach
  • OFA network training은 multi-objective problem으로 볼 수 있음
    1. sub-network 반복마다 OFA를 전체 objective function의 gradient로 최적화하면?
       -> sub-network 개수에 따라서 선형적으로 cost 증가
    2. 모두 사용하지 않고 일부 sub-network만 sampling하면?
       -> 많은 sub-network가 가중치를 공유하기 때문에 이 방법은 sub-network 정확도 유지에 방해

-> 그러니까 Progressive Shrinking 쓰자

 

• Progressive Shrinking
  • sub-network의 간섭을 방지하기 위해서, training 순서를 큰 sub-network에서 작은 sub-network로 수행하는 것
  • Progressive Shrinking
    1. Max kernel size, Max depth, Max width 가지는 large network training
    2. 작은 sub-network를 sampling space에 추가해서 작은 sub-network도 지원하도록 fine tune
    3. 제일 큰 network를 training하고 난 후, elastic kernel size를 선택하도록 함
       -> depth, width는 Max 고정
    4. 이후 elastic depth, elastic width 선택
       -> elastic resolution은 전체 training process에서 고정하지 않음
       -> training data의 각 batch에서 서로 다른 이미지 사이즈를 sampling하는 방식으로 적용
    5. (extra) 가장 큰 network 학습한 이후 Knowledge Distillation 적용
       -> 가장 큰 network로 얻은 soft label와 real label로 구성된 loss 활용
  • vs. Naive Approach
    1. OFA network가 작은 sub-network를 지원하도록 fine tune 된 경우
       -> progressive shrinking은 큰 sub-network가 well-trained되어 있으므로 작은 sub-network가 큰 sub-network를 방해하는 것을 방지할 수 있음
    2. 작은 sub-network와 큰 sub-network의 가중치 공유
       -> well-trained 된 큰 sub-network의 중요한 가중치를 가지는 작은 sub-network로 initialize 가능
       -> training process 빨라짐

Progressive Shirinking

• Elastic Kernel Size, Elastic Depth, Elastic Width
  1. Elastic Kernel Size
     • $7\times7$ 중간에 $5\times5$ 중간에 $3\times3$ 배치
     • 가운데의 $5\times5$, $3\times3$은 여러가지 역할을 수행할 수 있어야 함
       -> 독립적인 커널, 큰 커널의 일부
       -> 가운데의 sub-kernel들은 여러가지 역할에 대한 서로 다른 분포, 크기를 가져야 함
       -> 만약 강제로 동일하게 만들면 sub-network의 성능이 저하
     • Kernel Transformation Matrix
       -> kernel 가중치를 공유하기 위해서 사용
  2. Elastic Depth
     • N개의 layer를 가지는 unit에서 D개의 layer를 가지는 sub-network를 얻기 위해서, 첫번째 D개의 layer는 유지하고 마지막 N-D개의 layer는 건너뜀
     • 첫번째 D개의 layer에 대한 가중치는 큰 모델, 작은 모델 모두에게 공유
  3. Elastic Width
     • 각 layer에서 서로 다른 channel expansion ratio를 선택하도록 함
     • Channel Sorting Operation
       -> Progressive Shrinking을 따라서 full-width를 학습하고, partial width를 지원
       -> channel 가중치에 대한 L1 Norm을 이용해서 중요도를 판단하는 방식
       -> 큰 L1 Norm이면 중요한 channel인 것
       
       

Elastic Kernel

Elastic Width

- Specialized Model Deployment with OFA network

• OFA 학습 이후 주어진 시나리오에 특화된 sub-network를 찾는 단계
• 목표 : target hardware의 efficiency 제약(latency, energy)을 만족하면서 accuracy가 좋은 네트워크를 찾아내는 것
  • OFA는 training, search를 분리하기 때문에 이 단계에서는 training cost가 없음
  • neural network twins를 통해 주어진 network architecture에 대한 latency, accuracy 예측
    -> 모델 성능에 대한 빠른 feedback을 위함
• For example,
  1. 16K sub-network sampling
  2. 10K validation image에 대한 accuracy 측정
  3. [architecture, accuracy] pair를 이용해서 accuracy predictor를 학습
  4. 각 hardware에 대한 latency lookup table 구해서 latency predict
  5. 주어진 target hardware, latency 제약조건에서 neural-network-twins를 기반으로 Evolutionary Search

3. Experiments

- Settings

• dataset : ImageNet
• device : Samsung S7 Edge, Note 8, Note 10 / Google Pixel 1, Pixel 2 / LG G8 / Nvidia 1080Ti, V100 GPU / Jetson TX2, Intel Xeon CPU / Xilinx ZU9EG, ZU3EG FPGA

- Training OFA network on ImageNet

Progressive Shirinking 적용 결과 비교

• MobileNetV3과 동일한 architecture space 사용
• Progressive Shrinking 적용하면 sub-network의 ImageNet 정확도를 향상시킬 수 있음

- Specialized Sub-networks for Different Hardware and Constraints

OFA 성능 비교

CO2 배출량

• OFA는 Cost가 일정하지만 다른 NAS 방법은 deployment 시나리오 N에 대해 선형적임
• CO2 배출량을 보면, 다른 NAS와 비교했을 때 OFA의 효율성이 더 두드러짐

OFA vs. EfficientNet

• OFA가 hardware에서 EfficientNet 보다 빠르게 실행됨
• Search된 neural architecture를 처음부터 훈련해도 OFA와 동일한 성능을 내지 못함
  • architecture 뿐만 아니라 pre-train된 가중치도 OFA 성능에 영향을 줌

OFA vs. MobileNetV3

• OFA는 한번만 train 해도 좋은 성능을 보임 (초록색 선)
  • 이전의 NAS는 training cost가 많이 들기 때문에 불가능

Arithmetic Intensity 비교

Hardware 가속환경에서의 결과

• OFA를 특수 Hardware에 적용하는 경우 (FPGA, GPU)
  • 원칙 1 : 효율적인 CNN은 소량의 메모리 공간 만을 활용해서 여러 작업을 수행해야 함
    - Evolutionary Search는 더 좋은 연산을 가지는 CNN architecture를 찾을 수 있음
    - accuracy/latency feedback을 제공하는 OFA model twin 덕분
  • 원칙 2 : CNN은 hardware 가속기의 cost model과 함께 설계되어야 함
    - FPGA 가속기에는 point-wise 엔진이 있는데 $3\times3$ kernel과 동일
    - FPGA에서는 search space에 $5\times5$, $7\times7$이 있어도 탐색되지 않음
    - CPU에서는 대부분 대형 kernel이 자주 사용