Neural Architecture Search (NAS) - Once-for-All

https://github.com/mit-han-lab/once-for-all

GitHub - mit-han-lab/once-for-all: [ICLR 2020] Once for All: Train One Network and Specialize it for Efficient Deployment

 이번에 소개할 방법은 Once-for-all 이라는 방법이다. 2020년 ICLR에 나온 논문이고 논문 제목은 Once-For-All: Train one network and specialize it for efficient deployment 이다. 즉, 하나의 네트워크를 훈련시키고 그 네트워크 안에서 여러 서브 네트워크를 만들 수 있게 해서 각 상황에 맞는 서브 네트워크로 사용하겠다는 방법이다. 아래 그림처럼 스마트폰 별로 네트워크를 다르게 변형해서 줄 수도 있고 배터리별, cpu별, gpu별 등등 다양한 제약 상황에서 그에 맞는 네트워크를 하나의 네트워크로 부터 변형시켜 주는 것이다. 

 

 이러한 네트워크를 어떻게 만들고 학습시키는지 간단히 보도록 하겠다. 

 

 우리가 흔히 아는 CNN 네트워크에서 조절할 수 있는 변수들은 kernel size, channel수, layer 수, input 이미지의 resolution 등이 있다. 먼저 가장 큰 네트워크를 하나 학습시키고 여기서 이런 변수의 값들을 줄여서 작은 네트워크들을 만들어내게 된다. 그냥 단순히 하나의 네트워크를 학습시키고 이런 변수를 조절해 서브 네트워크를 뽑아내면 성능이 매우 떨어지게 된다. 이 논문에서는 progressive shrinking 방법으로 작은 네트워크를 뽑아내게 해서 이런 성능 하락을 최소화했다고 한다. 대략적인 방법은 아래 ppt에 나와있다.

 

 일단 여기서 조절하는 네트워크의 변수는 kernel size, layer수 channel수이다. 먼저 5개의 unit을 사용하고 각 unit에서 layer수는 [2,3,4] 중에 하나로 사용할 수 있다. 각 layer의 channel은 [3,4,6] 의 expansion ratio를 사용한다. 또한 kernel size는 [3,5,7] 중에 하나를 사용한다. 이렇게 하면 만들어질 수 있는 서브 네트워크의 수가 거의 10의 19승이 넘는다. 이런 수많은 서브 네트워크를 일일이 만들고 학습시키는 건 불가능한데 이렇게 하나의 네트워크만 학습시키면 그 수많은 네트워크를 얻어낼 수 있다는 것이기 때문에 굉장히 효율적인 방식이 된다. 

 

 먼저 kernel을 조절하는 방법을 보자. 처음에 가장 큰 size로 학습이 되기 때문에 7x7 커널을 갖고 있다. 이후에 이 커널을 5x5 혹은 3x3으로 줄여야 하는데 이 때 단순히 그냥 7x7에 있는 weight들을 그대로 가지고 5x5이나 3x3을 만들면 성능이 매우 떨어지게 돼서 transform matrix를 하나 추가해서 얻어낸다.  

 그리고 layer를 조절하는 방법을 보면 4개의 layer에서 줄일 때는 가장 마지막 단에 있는 layer부터 사용하지 않으면서 layer수를 줄인다. 

 다음으론 channel 수를 조절하는 방법을 보면 각 channel의 importance를 측정해 중요도가 낮은 순서대로 사용하지 않는데 여기서 중요도는 L1 norm 즉, 각 weight의 크기대로 중요도를 매긴다. 

 순서는 위의 그림대로 kernel부터 조절하면서 학습시키고(layer수와 channel수는 최대로 유지하면서) 다음으론 layer수, channel수를 줄여가면서 학습을 시킨다. 전체 학습과정에서 input 이미지의 resolution도 계속 다르게 넣어서 다양한 resolution에도 결과가 잘 나오게 한다. 이런 progressive shrinking 방식으로 변수를 조절해 아래 그림에 나와있는 것처럼 네트워크를 줄이고 학습시키고 줄이고 학습시키고를 반복해서 최종 네트워크를 얻어낸다. 마치 pruning을 하고 fine-tuning을 하는 과정과 비슷하지만 다양한 dimension에서 네트워크를 줄인다는 점이 다르다. 

  

 그리고 특정 hardware의 constraints가 들어오면 그에 맞는 네트워크를 만들어 내야하기 때문에 architecture search과정에서 accuracy predictor와 latency predictor또한 같이 학습시킨다. 이 predictor를 가지고 나중에 특정 constraint가 들어오면 그에 맞는 네트워크를 만들어낼 수 있도록 한다.