如何使用MindSpore自定义优化器

引言

    神经网络的参数众多，我们需要选择合适的算法来进行参数的更新和学习，也就是优化器。优化器在神经网络模型训练的过程中有着十分重要的作用。
    从SGD开始，神经网络模型优化器就一直在迭代和发展之中。如PyTorch就已经开源了包括SGD、Momentum、RMSprop、Adam、AdamW等等丰富的优化器。但是，由于深度学习模型本身的复杂性，深度学习模型框架自带的优化器本身可能并不能很好的适应我们的任务需求，因此我们有时候需要根据自己的任务去定义合适的优化器，让深度学习模型的优化能更加顺利。

动机

    要写这篇关于《如何使用MindSpore优化器》博客的起因是由于在复现DeppMind的论文《High-Performance Large-Scale Image Recognition Without Normalization》，也就是NFNet。关于NFNet的介绍，大家可以参考博客最强ResNet变体！归一化再见！DeepMind提出NFNet，代码已开源！。

层归一化的缺陷

    在计算机视觉领域，神经网络的训练过程中，虽然诸如BatchNorm（主要用在CNN中，一下称为BN）、LayerNorm（主要用在ViT中，以下称为LN）能够为训练带来很稳定的收敛，在一定程度上可以提高模型的指标，但是两者带来的弊端也是不容忽视的。对于BN层来说，模型最后的精度很容易受到BatchSize的影响，当BatchSize很小的时候，模型最后的指标会发生严重的下降。为了缓解这一现象，研究者们提出了很多替换BN层的归一化方法**，目前在ViT中广泛使用的LN**。尽管LN具有样本之间的独立性，模型的最后指标几乎不会受到BatchSize的影响，但是LN层的存在却会大大的降低模型的单步训练时长，也可以说：LN本身的一个硬件不友好的operation。

自适应梯度裁剪

    从作者的论文中可以看到，作者使用了一种名为自适应梯度裁剪的技术，相关的代码可以参见nfnets_pytorch。这里我将结合代码和数学公式对论文里面的自适应梯度裁剪技术进行说明。

"""
unitwise_norm的功能:
    对于Scalars等只有一个维度的张量：求出其二范数
    对于全连接层的权重：对于OxI的矩阵，求出每一个I的2-范数
    对于卷积层的权重：OIHW的矩阵，求出每一个IHW的F范数（其实也就是flatten的2-范数）
"""
def unitwise_norm(x):
    if (len(torch.squeeze(x).shape)) <= 1: # Scalars, vectors
    	# NFNet中的缩放因子，squeeze之后实际上只有一个维度
        axis = 0
        keepdims = False
    elif len(x.shape) in [2,3]: # Linear layers
        # Original code: IO
        # Pytorch: OI
        axis = 1
        keepdims = True
    elif len(x.shape) == 4: # Conv kernels
        # Original code: HWIO
        # Pytorch: OIHW
        axis = [1, 2, 3]
        keepdims = True
    else:
        raise ValueError(f'Got a parameter with len(shape) not in [1, 2, 3, 4]! {x}')

    return torch.sqrt(torch.sum(torch.square(x), axis=axis, keepdim=keepdims))

d_p = p.grad

# =========================
# Gradient clipping
if clipping is not None:
	"""对于训练权重p，按照unitwise_norm函数求出对应权重的2-范数，限定eps为最小阈值"""
    param_norm = torch.maximum(unitwise_norm(p), torch.tensor(eps).to(p.device))
    """对于训练权重p的梯度，按照unitwise_norm函数求出对应梯度的2-范数，这里不加阈值"""
    grad_norm = unitwise_norm(d_p)
    """这里给定一个 clipping 缩放系数 """
    max_norm = param_norm * group['clipping']
    
  	"""梯度的对应的范数 和 权重对应的缩放范数"""
    trigger_mask = grad_norm > max_norm
    """截断的梯度= 梯度 * 权重对应的缩放函数 / max(梯度范数, 最小阈值)"""
    clipped_grad = p.grad * (max_norm / torch.maximum(grad_norm, torch.tensor(1e-6).to(p.device)))
    """
    最后输出梯度: 
    	如果 梯度范数 > 缩放后的权重范数
    		梯度 = 截断梯度
    	否则
    	 	梯度 = 本身不变
    """
    d_p = torch.where(trigger_mask, clipped_grad, d_p)
    """这里猜测，加入最小一直可能只是为了防止最后的输出接近0，导致出现一些特别大或者特别小的数字，
    在理解具具体公式的时候，我们可以暂时不管这两个eps"""

$$G\to \{\begin{array}{rcl}\lambda \frac{||W_i^{\iota }|{|}_F^{\ast }}{||G_i^{\iota }|{|}_F}{G}_i^{\iota }& & \frac{||W_i^{\iota }|{|}_F^{\ast }}{||G_i^{\iota }|{|}_F}>\lambda \\ G& & otherwise\end{array}$$

    通过理解公式我们可以简单理解为，当梯度的范数大于一定程度的权重范数时，梯度会以
$$\frac{weight\_norm}{grad\_norm}<1$$进行缩放，反之梯度保持为原来的值，这就是其中自适应三个字的由来。

如何用MindSpore自定义优化器并且实现AGC_SGD

    要用MindSpore优雅的实现AGC_SGD（主要是想尽量少用for循环，尽量接近MindSpore 的原生程序风格），我们首先要来了解一下MindSpore优化器的整体优化器的优化构建。以下的代码片段有一些跳跃，希望大家紧跟思路，抓住核心的代码片段就好。

mindspore.nn.optim.Momentum

    MindSpore的优化器整体定义和其他框架类似，都是从一个Optimizer基类继承来的op。在PyTorch中SGD似乎是和Momentum写在一起的，因此这里拿Momentum为例。对于Momentum而言，

变量名	含义
grad	gradients（梯度）
lr	learning_rate（学习率）
p	params（模型参数）
v	moments（梯度的动量）
u	momentum（动量系数）

以下的Momentum优化器的更新准则（暂时忽略weight_decay）
首先是得到
$$当前梯度动量=过去动量\times 动量系数+（1-动量系数）\times 当前的梯度$$
用公式表示也就是：
$$v_t=v_{t-1}\ast u+(1-u)\times gradients$$
当use_nesterov=False的时候，就完成单步的更新，同时，相对应的动量也会得到更新。关于use_nesterov参数，感兴趣的读者可以查找牛顿动量（Nesterov）算法
If use_nesterov is True:

$$p_t=p_{t-1}-(grad\times lr+v_t\times u\times lr)$$
If use_nesterov is False:
$$p_t=p_{t-1}-lr\times v_t$$
可以看到，在更新的过程中，我们既可以等到为下一步准备的当前动量，也可以完成最后参数的更新，以此往复，不断更新参数。

class Momentum(Optimizer):
    def __init__(self, params, learning_rate, momentum, weight_decay=0.0, loss_scale=1.0, use_nesterov=False):
        super(Momentum, self).__init__(learning_rate, params, weight_decay, loss_scale)
        Validator.check_value_type("momentum", momentum, [float], self.cls_name)
        if isinstance(momentum, float) and momentum < 0.0:
            raise ValueError("momentum should be at least 0.0, but got momentum {}".format(momentum))
        self.momentum = Parameter(Tensor(momentum, mstype.float32), name="momentum")
        self.params = self.parameters
        self.use_nesterov = Validator.check_bool(use_nesterov)
        # 复制一份和self.params一样形状的参数，初始化为0，用来作为保存动量的副本
        self.moments = self.params.clone(prefix="moments", init='zeros')
        self.hyper_map = C.HyperMap()
        # 动量优化器，实施算法
        self.opt = _selected_ops.ApplyMomentum(use_nesterov=self.use_nesterov)

    在优化器的运行阶段，我们可以看到，大体上就是执行了

1. 权重衰减
2. 梯度缩放（估计是为了配和缩放器加的参数，假定scale=1.，我们传入优化器前就对混合精度完成缩放，可以看到我之前的如何使用MindSpore进行自定义训练）
3. 梯度中心化（这个功能默认是False，op里面不会去执行，因此不用太过于在意）
4. 完成单步更新（group_lr是指针对不同的参数使用不同的lr，感兴趣的可以尝试）

        params = self.params
        moments = self.moments
        gradients = self.decay_weight(gradients)
        gradients = self.scale_grad(gradients)
        gradients = self.gradients_centralization(gradients)
        lr = self.get_lr()
        if self.is_group_lr:
            success = self.hyper_map(F.partial(_momentum_opt, self.opt, self.momentum), lr, gradients, params, moments,
                                     self.ps_parameters, self.cache_enable)
        else:
            success = self.hyper_map(F.partial(_momentum_opt, self.opt, self.momentum, lr), gradients, params, moments,
                                     self.ps_parameters, self.cache_enable)
        return success

    在Momentum中，我们可以看到一个if的逻辑判断，从大体上就可以了解到，估计就是判断是不是优化器中包含过往的动量信息，_ps_pull和_ps_push操作估计就是为了保存下过去的缓存信息。

_momentum_opt = C.MultitypeFuncGraph("momentum_opt")

@_momentum_opt.register("Function", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Bool", "Bool")
def _tensor_run_opt_ext(opt, momentum, learning_rate, gradient, weight, moment, ps_parameter, cache_enable):
    """Apply momentum optimizer to the weight parameter using Tensor."""
    if ps_parameter and not cache_enable:
        op_shape = P.Shape()
        _ps_pull = P.Pull()
        _ps_push = P.Push("ApplyMomentum", [])
        shapes = (op_shape(learning_rate), op_shape(gradient), op_shape(momentum))
        success = F.depend(True, _ps_pull(_ps_push((learning_rate, gradient, momentum), shapes), weight))
    else:
        success = F.depend(True, opt(weight, moment, learning_rate, gradient, momentum))
    return success

使用MindSpore实现AGC

    经过上面的介绍，我们可以了解到，实际上在MindSpore实现AGC_SGD，只要在Momentum里面增加一个AGC操作就可以了。以下是代码，这里就不过多赘述了。
    需要注意的是，在MindSpore中，静态图运行下对于数值的操作都是需要使用C.MultitypeFuncGraph函数构建功能图的。关于这个Op大致就是构建了一个可以用在map函数里面的一个function，然后对变量实行统一操作。目前这个op里面如果包含if的话，最后一定要配上else这种显式的完备逻辑，否则程序会报错的。（就是和PyTorch不咋一样的原因，功能上不影响）

_agc_clip = C.MultitypeFuncGraph("agc_clip")


@_agc_clip.register("Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Bool")
def _tensor_run_agc_ext(eps, min_grad, clipping, grad_norm, gradient, weight_norm, clip):
    """Apply sgd optimizer to the weight parameter using Tensor."""
    if clip:
        param_norm = ops.Maximum()(weight_norm, eps)
        max_norm = param_norm * clipping
        trigger_mask = ops.Greater()(grad_norm, max_norm)
        clipped_grad = gradient * (max_norm / ops.Maximum()(grad_norm, min_grad))
        gradient = mnp.where(trigger_mask, clipped_grad, gradient)
    return gradient


_unitwise_norm = C.MultitypeFuncGraph("unitwise_norm")


@_unitwise_norm.register("Tensor")
def unitwise_norm_solve(x):
    if (len(ops.Squeeze()(x).shape)) <= 1:  # Scalars, vectors
        axis = 0
        keepdims = False
    elif len(x.shape) in [2, 3]:  # Linear layers
        # Original code: IO
        # Pytorch: OI
        axis = 1
        keepdims = True
    else:
        # Conv kernels
        # Original code: HWIO
        # Pytorch: OIHW
        axis = (1, 2, 3)
        keepdims = True
    return ops.Sqrt()(ops.ReduceSum(keepdims)(ops.Square()(x), axis))

    另外，根据作者的意思，最后一层的全连接层的weight权重是不需要使用AGC为好，因此我们可以构建 self.group_clipping_tuple: List[Bool,]对其进行针对name的判断。在优化器中，为了针对不同的权重(weight、bias、gamma、beta），针对权重衰减的weight_decay也会根据参数的顺序转化成一个列表进行权重和weight_decay的一一映射关系。具体的代码可以参见Optimizer源码

        self.group_clipping_tuple = []
        for index, params in enumerate(self.group_params):
            name = params.name
            if not "attn_last" in name and "fc" in name and 'bias' not in name:
                print(f"{name} no clipping")
                self.group_clipping_tuple.append(False)
            else:
                self.group_clipping_tuple.append(True)
        self.group_clipping_tuple = tuple(self.group_clipping_tuple)
        assert len(self.group_clipping_tuple) == self.param_length
       

    至此，我们就依赖于MindSpore本身自带的Momentum函数，完成了MindSpore实现AGC_SGD的功能，代码过段时间我会提交MindSpore官方的model仓库，刚兴趣的小伙伴到时候可以自取。

MindSpore社区贡献活动

    这里帮小老哥再打个小广告，MindSpore社区贡献活动，主要就是使用MindSpore复现顶会论文，既能掌握MindSpore框架的使用，为MindSpore开源社区建设出力，还可以收获超值奖品。活动应该会持续到2021年年底，大家快来参加吧！
活动链接：https://bbs.huaweicloud.com/forum/thread-86967-1-1.html