Accurate, Large Minibatch SGD - Training ImageNet in 1 Hour
by Yeonjee Jung
[Abstract]
딥러닝에서 네트워크와 데이터셋이 커지면 학습시간이 늘어나게 되어서 분산병렬 SGD를 사용한다. 이 논문에서는 미니배치 크기가 크면 최적화하는 데에 어려움을 야기하지만, 이렇게 훈련된 네트워크는 일반화를 잘 한다는 것을 보여준다. 또한 한번에 미니배치 사이즈를 8192개까지 늘려도 정확도에는 변함이 없다는 것을 보여준다. 이 결과를 얻기 위해 learning rate를 정하는 규칙을 적용했고, 새로운 warmup 방법을 만들었다.
[2] Large Minibatch SGD
[2.1] Learning Rates for Large Minibatches
이 논문의 목적은 정확도는 유지하면서 큰 미니배치 사이즈를 사용하는 것이다. 이들은 Linear Scaling Rule이라고 불리는, 미니배치 크기가 $k$배 커지면 learning rate도 $k$배 크게 해야 한다는 규칙을 발견했다. 작은 미니배치를 사용했을 때의 $w_{t+k}$와, $k$배의 미니배치를 사용했을 때의 $\hat{w}_{t+1}$이 비슷해지려면 $\hat{\eta}=k\eta$가 되어야 한다. 이렇게 학습시켰을 때, 둘의 정확도와 학습곡선이 비슷하다. 그러나 학습 초기에는 비슷하지 않을 수 있고, 미니배치 사이즈는 계속해서 커질 수 없다는 성질이 있다. 이 논문은 전례 없는 미니배치 사이즈를 가지고 실험적으로 Discussion에 있는 이론들을 테스트해본다.
[2.2] Warmup
constant warmup은 학습 초기 epoch에서 constant learning rate를 사용한다. 그러나 이 논문에서는 충분하지 않아서 gradual warmup을 사용했는데, 훈련 초반에 잘 수렴하도록 해준다.
[2.3] Batch Normalization with Large Minibatches
BN을 사용하면 원래 집합 $X$안에 있는 크기가 $n$인 모든 부분집합들을 모두 포함하고 있는 것처럼 생각할 수 있다. 분산 컴퓨팅 환경에서 BN을 사용한다면 per-worker sample size가 $n$이고 total minibatch size가 $kn$일 때, $X^n$에서 골라낸 독립된 $k$개의 $\cal{B}_ j$로 볼 수 있다. 이게 결국 각 minibatch의 분산을 비슷하게 만들어준다는 수업때 하려던 내용인듯
[3] Subtleties and Pitfalls of Distributed SGD
Weight decay
L2 regularization때문에 weight decay를 사용하는데, 이 항이 원래의 loss에 더해진다. 따라서 원래 loss를 scaling 하는 것과 learning rate를 scaling하는 것은 다르다.
Momentum correction
momentum을 사용한 SGD가 최근 많이 이용되고 있는데, learning rate가 바뀌면 momentum correction term을 적용해야 한다. 이게 왜 문제점인지 잘 모르겠음
Gradient aggregation
전체 loss를 계산할 때 평균을 사용하는데, 각 worker의 loss를 $n$으로 나누면 다시 한번 전체를 $k$로 나누어야 하기 때문에, 그냥 처음부터 각 worker의 loss를 $kn$으로 나눠서 최종적으로는 다 더하기만 한다.
Data shuffling
한 epoch당 데이터 전체를 $k$ worker에게 분배하는데, 모든 데이터들은 반드시 한번씩만 사용되도록 분배한다.
[4] Communication
gradient aggregation은 backprop과 병렬로 처리되어야 한다. 한 layer의 gradient가 계산되면, 모든 worker에게서 정보를 받아서 합치고, 그와 동시에 다음 layer의 gradient가 계산될 수 있다.
[4.1] Gradient Aggregation
gradient aggregation에는 allreduce 연산이 사용된다. 이 논문에서의 구현은 세 부분으로 되어 있다.
- 8개의 GPU에 있는 버퍼들을 각 서버에서 하나로 합친다.
- 결과 버퍼가 공유되고 합쳐진다.
- 결과가 각 GPU로 퍼진다.
1번과 3번은 NCCL을 이용하였다. 2번에 대해서는 recursive halving and doubling 알고리즘과 bucket 알고리즘을 구현했는데, 큰 사이즈의 요소에 대해서는 recursive halving and doubling 알고리즘이 더 좋았다.
having/doubling 알고리즘은 reduce scatter collective와 allgather 단계로 구성된다. reduce scatter에서는 서버가 2개씩 짝을 지어, 1번 서버의 경우 자신의 버퍼의 뒤쪽 반을 보내고 2번 서버의 앞쪽 반을 받는다. 따라서 2의 승수개의 서버가 필요하다. 수신 및 전송하는 데이터가 reduction 되는 동안 목적지까지의 거리는 2배가 된다. 이 단계 이후에는 각 서버에는 최종 reduction된 데이터의 일부가 있다. allgather 단계에서는 통신을 역추적해서 각 서버에 있는 벡터들을 모은다. 각 서버에서 처음에 보내는 부분이 allgather에서 수신되고 있고, 그 이후에는 받았던 부분을 보낸다.
[4.2] Software
이부분이 사실 메인인 것 같은데.. 이해가 안감ㅠㅠ
[4.3] Hardware
이 논문에서는 50Gbit 네트워크가 ResNet-50을 학습시키는 데에 충분하다고 한다. 게다가 이 최대 bandwidth는 backprop에만 사용되고, forward pass에는 그만큼 쓰이지 않는다.
[5] Main Results and Analysis
이 논문의 메인 결과는 ImageNet에 ResNet-50과 256worker를 이용해 정확도 손실 없이 한시간만에 학습시켰다는 것이다. Linear scaling rule과 warmup 전략을 사용해서 이가 가능했다.
[5.2] Optimization or Generalization Issues?
이 논문에서는 no warmup, constant warmup, gradual warmup을 실험했는데, training error와 validation error 면에서 모두 gradual warmup이 가장 성능이 좋았다. 또한 baseline의 learning curve도 초반 빼고는 잘 따라갔다.
[5.3] Analysis Experiments
minibatch size가 8k보다 작거나 같을 때에는 learning curve가 거의 비슷했는데, 그 이상에서는 비슷하지도 않았다. 여기서 알 수 있는 것은, learning curve를 비교해 보면 training이 끝나기 전에 결과가 잘 나올지, 아닐지 알 수 있다는 점이다.
minibatch size가 작았을 때는 $\eta=0.1$언저리에서는 비슷한 성능이 나왔지만, size가 커질수록 $\eta$를 더 섬세하게 바꿔야 한다.
BN에서의 일반적인 초기화는 $\gamma=1$인데, 이 논문에서는 각 residual block에 대해서만 $\gamma=0$으로 초기화했더니 더 좋은 결과가 나왔다.
ImageNet-5k dataset에 대해서도 실험해봤는데, 정확도가 작은 minibatch size와 정확도가 비슷한 minibatch size는 8k까지로, 데이터셋의 크기와 minibatch size의 관계는 없는 것으로 관찰되었다.
[5.4] Generalization to Detection and Segmentation
ImageNet을 학습시킨 ResNet-50을 R-CNN의 초기화에 이용해보았는데, $8k$까지는 성능이 나빠지지 않았다. linear scaling rule과 large batchsize는 다른 모델에서도 잘 적용된다는 것을 보여준다.
[5.5] Run Time
batch size를 늘리면 GPU간의 소통에도 불구하고 iteration 당 시간은 비슷하지만, epoch당 걸리는 시간이 선형적으로 줄어든다. 또한 GPU가 k배 늘어나면 초당 처리되는 이미지 수도 k배 늘어나는 것이 이상적인데, 이 논문의 방법은 GPU간의 소통에도 불구하고 이상의 $90%$정도를 달성했다.
allreduce operation
Subscribe via RSS