Scaling SGD Batch Size to 32K for ImageNet Training

为了充分利用GPU计算,加快训练速度,通常采取的方法是增大batch size.然而增大batch size的同时,又要保证精度不下降,目前的state of the art 方法是等比例与batch size增加学习率,并采Sqrt Scaling Rule,Linear Scaling Rule,Warmup Schem等策略来更新学来率.

在训练过程中,通过控制学习率,便可以在训练的时候采用大的batch,例如,batch=1024 for AlexNet,batch=8192 for ResNet-50.

那么,如何当有多个gpu的时候,如何利用更大的batch size.本文没有使用更大的学习率,而是提出了一种 Layer-wise Adaptive Rate Scaling (LARS)方法,LARS LR方法是,对不同的层采用不同的LR(学习率),每层的LR的值有该层的weigthts,gradients相关.通过使用LARS LR算法,k可以增大batch size,32768 for ResNet,8192 for AlexNet.

LARS LR算法定义如下:

$\eta=l\times\gamma\times\frac{||w||_2}{||\nabla w||_2}$

式中,$l$为scaling factor,在AlexNet,ResNet训练中,$l=0.001$,$\gamma$为输入学习率LR,$\gamma$取值范围通常为1到50.$w$为每层的weight,$\nabla w$为weight的梯度.

更新weights:

$w=w-\eta\nabla w$