CUDA编程入门之优化器Momentum_世纪娱乐-世纪注册站

CUDA编程入门之优化器Momentum

公司相册 | 2024-04-07 22:57

上一篇主要介绍了经典的梯度下降法算法并阐述了其存在的一些局限，例如，在训练过程中，当接近最优值时梯度会比较小，由于学习率固定，普通的梯度下降法的收敛速度会变慢，有时甚至陷入局部最优。要想到达最优解，需要我们不断的迭代或者调整学习率，但是调大学习率会导致每一次迭代的步长过大，也就是摆动过大，误差较大；调小学习率会让迭代次数增加，从而明显增加训练时间。这时如果考虑历史梯度，将会引导参数朝着最优值更快收敛，这就是动量算法的基本思想。

动量优化方法是在梯度下降法的基础上进行的改变，具有加速梯度下降的作用。一般有标准动量优化方法Momentum、NAG（Nesterov accelerated gradient）动量优化方法。

使用动量(Momentum)的随机梯度下降法(SGD)，主要思想是引入一个积攒历史梯度信息动量来加速SGD；假设从训练集中取一个大小为 $n$ 的小批量样本 ${X^{1}, X^{2},..., X^{n}}$ ，对应真实值为 $Y^{i}$ ，Momentum优化表达式表示如下：

$W_{t+1}=W_{t}- v_{t}，$

$v_{t}=\\alpha v_{t-1}+ \\eta_{t}\\Delta J(W_{t},X^{i_{s}}, Y^{i^{s}})$

其中 $v_{t}$ 表示 $t$ 时刻积攒的加速度， $\\alpha$ 表示动力的大小，一般取值为0.9（表示最大速度10倍于SGD）， $\\Delta J(W_{t},X^{i_{s}}, Y^{i^{s}})$ 表示随机选择的一个梯度方向， $W_{t}$ 表示 $t$ 时刻模型参数， $\\eta_{t}$ 表示学习率；

动量主要解决SGD的两个问题：一是随机梯度的方法（引入的噪声）；二是Hessian矩阵病态问题（可以理解为SGD在收敛过程中和正确梯度相比来回摆动比较大的问题）

这个策略也可以这样理解，当我们将一个小球从山上滚下来时，没有阻力的话，它的动量会越来越大，但是如果遇到了阻力，速度就会变小。加入的这一项，可以使得梯度方向不变的维度上速度变快，梯度方向有所改变的维度上的更新速度变慢，这样就可以加快收敛并减小震荡。但这种情况相当于小球从山上滚下来时是在盲目地沿着坡滚，如果它能具备一些先知，例如快要上坡时，就知道需要减速了的话，适应性会更好。

牛顿加速梯度（NAG, Nesterov accelerated gradient）算法，是Momentum动量算法的变种。更新模型参数表达式如下：

$W_{t+1}=W_{t}- v_{t}，$

$v_{t}=\\alpha v_{t-1}+ \\eta_{t}\\Delta J(W_{t}-\\alpha v_{t-1})$

其中 $v_{t}$ 表示 $t$ 时刻积攒的加速度， $\\alpha$ 表示动力的大小， $W_{t}$ 表示 $t$ 时刻模型参数， $\\eta_{t}$ 表示学习率， $\\Delta J(W_{t}-\\alpha v_{t-1})$ 表示损失函数关于 $W_{t}$ 的梯度；

Nesterov动量梯度的计算在模型参数施加当前速度之后，因此可以理解为往标准动量中添加了一个校正因子；

这个策略也可以这样理解，在Momentun中小球会盲目地跟从下坡的梯度，容易发生错误。所以需要一个更聪明的小球，能提前知道它要去哪里，还要知道走到坡底的时候速度慢下来而不是又冲上另一个坡。计算 $\\Delta J(W_{t}-\\alpha v_{t-1})$ 可以表示小球下一个位置大概在哪里，从而可以提前知道下一个位置的梯度，然后使用到当前位置来更新参数。

一般深度学习框架都会有 Momentum 算法的实现，此处来看下 Caffe2 中关于 MomentumSGD 的 CUDA 实现。

Momentum

template <>
__global__ void MomentumSGDKernel<false>(
    const int N,
    const float* g,
    const float* m,
    float* ng,
    float* nm,
    const float* lr,
    const float momentum,
    float* param) {
  const float LR = lr[0];
  CUDA_1D_KERNEL_LOOP(i, N) {
    const float adjusted_gradient = LR * g[i] + momentum * m[i];
    nm[i] = adjusted_gradient;
    ng[i] = adjusted_gradient;
    if (param != nullptr) {
      param[i] -= adjusted_gradient;
    }
  }
}

NAG

template <>
__global__ void MomentumSGDKernel<true>(
    const int N,
    const float* g,
    const float* m,
    float* ng,
    float* nm,
    const float* lr,
    const float momentum,
    float* param) {
  const float LR = lr[0];
  CUDA_1D_KERNEL_LOOP(i, N) {
    const float mi = m[i];
    const float mi_new = momentum * mi + LR * g[i];
    nm[i] = mi_new;
    ng[i] = fmaf(momentum, mi_new - mi, mi_new);
    if (param != nullptr) {
      param[i] -= ng[i];
    }
  }
}

此处 $N$ 表示权重大小，即前面提到的模型参数 $W$ 的 shape； $g$ 表示梯度，即前面提到的 $\\Delta J(W)$ ； $lr$ 表示学习率，即前面提到的 $\\eta_{t}$ ，momentum 表示动力的大小，即前面提到的 $\\alpha$ ，m[i]表示动量。nm 和 ng 用于保存梯度和动量，

这里再举一个mindspore的实现，对比学习；

template <typename T, typename S, typename G>
__global__ void MomentumUpdateVariableKernel(const size_t size, T *variable, T *accumulation, const S *learning_rate,
                                             const G *gradient, const S *momentum, bool use_nesterov) {
  if (use_nesterov) {
    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < (size); i += blockDim.x * gridDim.x) {
      accumulation[i] = momentum[0] * accumulation[i] + gradient[i];
      variable[i] -= gradient[i] * learning_rate[0] + accumulation[i] * momentum[0] * learning_rate[0];
    }
  } else {
    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < (size); i += blockDim.x * gridDim.x) {
      accumulation[i] = momentum[0] * accumulation[i] + gradient[i];
      variable[i] -= learning_rate[0] * accumulation[i];
    }
  }
}

Caffe2 还提供了另外一种实现，使用 float2 类型压缩存储向量，对于float2的概念可参考：

使用float2类型压缩存储向量_用户指南_云数据库RDS_敏捷版数据库场景

下面只列了 Momentum 实现，NAG 实现可参考前面的定义自己尝试去写一写；

__global__ void FP32MomentumSGDKernel(
    int N,
    const float2* g,
    const float2* m,
    float2* ng,
    float2* nm,
    const float* lr,
    const float mom,
    bool nesterov,
    const float wd,
    float2* param) {
  const float lr2 = lr[0];
  const float LR = lr2;
  const float momentum = mom;
  const float weight_decay = wd;

  int n = N / 2;
  CUDA_1D_KERNEL_LOOP(i, n) {
      ng[i].x = __fmaf_rn(weight_decay, param[i].x, g[i].x);
      ng[i].y = __fmaf_rn(weight_decay, param[i].y, g[i].y);

      float2 mi_float2 = m[i];
      float2 adjusted_gradient_float2;
      adjusted_gradient_float2.x =
          __fmaf_rn(LR, ng[i].x, __fmul_rn(momentum, mi_float2.x));
      adjusted_gradient_float2.y =
          __fmaf_rn(LR, ng[i].y, __fmul_rn(momentum, mi_float2.y));

      nm[i] = adjusted_gradient_float2;
      ng[i] = adjusted_gradient_float2;

      if (param) {
        param[i].x = __fsub_rn(param[i].x, adjusted_gradient_float2.x);
        param[i].y = __fsub_rn(param[i].y, adjusted_gradient_float2.y);
      }
    }
}

其中一些math 函数定义如下：

Compute $x\ imes y + z$ as a single operation, in round-to-nearest-even mode.

__device__? float __fmaf_rn ( float  x, float  y, float  z )

Compute the product of $x$ and $y$ in round-to-nearest-even mode.

__device__? float __fmul_rn ( float  x, float  y )

Compute the difference of $x$ and $y$ in round-to-nearest-even rounding mode.

__device__? float __fsub_rn ( float  x, float  y )