混合精度训练

Posted by KevinLT on September 14, 2018

tensor core介绍

tensor core是NVIDIA在其Volta架构显卡中引入的硬件单元,其目的在于加速矩阵运算以加快神经网络的训练速度。 tensor core的主要功能为进行混合精度的FMA(Fused Multiply-Add,乘积累加运算)。

tensor core运算过程
tensor core运算过程

单个tensor core单元能够在一个时钟周期内完成64次上图所示D = A X B + C运算。 其中A,B,C,D均为4X4的矩阵,A和B为FP16(半精度浮点数)矩阵,C和D可以为FP16或FP32(单精度浮点数)矩阵。 矩阵的运算是神经网络计算过程中最频繁的操作,tensor core通过将矩阵乘法与加法运算并行在了一起,大大加速了计算速度。

详细的Volta架构及tensor core介绍可以看这里

混合精度训练

上面我们有介绍到,在tensor core的FMA计算中,矩阵乘法是在FP16精度下进行的。 所以要使用tensor core来加速训练,网络参数需要以FP16精度表示。

FP16格式
FP16格式

FP16的格式如上图所示,其指数部分为5位,能表示范围在-14~15之间,小数部分为10位,故小于 2-24的量级FP16是无法表示的。

由于FP16的精度损失问题,如果我们在神经网络的训练过程中直接将网络参数以FP16的形式进行计算,可能会出现数值不稳定的情况而导致模型性能下降。 对此,百度和NVIDIA的研究院在Mixed Precision Training这一论文中提出混合精度训练的方法,在充分利用FP16加速运算的优点的同时保证了模型的精度。下面介绍其论文中的主要要点。

  1. FP32 master copy of weigths

    FP32 master copy即维护一份网络中FP16精度参数的FP32精度的拷贝。 计算过程如下图所示,在前向传播过程中,使用由master copy类型转换得到的FP16精度参数进行运算; 而在反向传播计算完梯度后,将梯度作用到master copy上以在FP32精度上进行参数更新。

    FP32 master copy
    FP32 master copy

    这么做的主要原因有两个。

    第一是由于若直接在FP16精度下进行参数更新,存在梯度过小而导致更新值为0的情况。 下图是某次模型训练过程汇总网络参数的梯度的分布直方图,有5%的梯度值是分布在小于2-24的区间内的。 若优化器直接将这部分梯度乘上学习率作用到FP16精度参数的更新上,那么更新值将为0。 这将影响到模型的准确率。 但是如果是更新值是作用到FP32精度的master copy上时,则不会出现更新值下溢为0的情况。

    第二是如果参数相较于其更新值过大的话,也可能会由于浮点数加法机制的缘故而导致更新值为0。 在浮点数加法过程中,需要将两数对齐进行运算。 如果参数的大小是其更新值的2048倍或者更大的话,那么更新值的小数位需要右移至少11位才能与前者对齐,这超出了FP16精度的表示范围。 在FP32精度下则一般不会出现这种问题。

    gradient histogram
    gradient histogram
  2. loss scaling

    loss scaling即将loss值放大,以保证反向传播过程当中梯度落在FP16精度能表示的范围之间。

    下图是Multibox SSD网络在训练过程过程中激活单元梯度的分布情况。其中有67%的梯度落在了小于2-24的范围内,在FP16精度下无法表示。 如果不对梯度进行放大,在FP16精度下对该网络进行训练将导致发散。 对激活梯度进行放大,再对参数梯度缩小相应倍数,即可解决该问题。

    activation gradients
    activation gradients

    根据梯度计算的链式法则,对梯度放大的最简单的方法就是放大loss。 放大因子的大小选择没有固定的标准,对于上述Multibox SSD网络,作者尝试了8-32K的放大因子,均训练成功了。 只要保证放大后的梯度不超过FP16精度的表示上限(65504),选择较大的放大因子并无副作用。

TensorFlow实现

下面的代码主要参考https://github.com/tensorflow/models/tree/master/official/resnet

  1. 将输入转为FP16精度

    inputs = tf.cast(inputs, tf.float16)
  2. 使用custom_getter构建模型实现FP32 master copy

    下面即是一个模型的样例代码。

    forward为模型的前向传播计算过程,我们使用一个_model_variable_scope来管理Variable上下文环境。 TensorFlow中的Vairable即网络中的参数。

    _model_variable_scope函数返回一个tf.variable_scope。 它是TensorFlow提供的定义Variables创建过程的上下文管理器,官方文档在这里。 第一个参数用于定义scope名称,所有在该上下文环境内创建的Variable均会加上该名称作为前缀。 custom_getter参数用于指定Variable获取函数。

    _custom_dtype_getter即我们设定的Variable获取函数。 在variable_scope中,默认的获取函数为tf.get_variable。 如果我们指定了custom_getter参数,那么tf.get_variable将调用我们设定的函数,而不是直接获取Variable。 custom_getter所指定的函数应当与tf.get_variable参数列表相同,除了它多出了一个getter参数。 该参数即为默认的tf.get_variable函数。 在该函数中,我们规定了如果要获取的Variable类型为tf.float16,那么我们首先获取其32位的master copy,再将其转换为16位返回。 其他类型的Variable则直接调用getter返回。 ``` class Model(object):

     def __init__(self, ..., dtype=tf.float16):
         self.dtype = dtype
         ...
    
     def _custom_dtype_getter(self, getter, name, shape=None, dtype=self.dtype,
                                *args, **kwargs):
         if dtype is tf.float16:
             var = getter(name, shape, tf.float32, *args, **kwargs)
             return tf.cast(var, dtype=dtype, name=name + '_cast')
         else:
             return getter(name, shape, dtype, *args, **kwargs)
    
     def _model_variable_scope(self):
         return tf.variable_scope('scope_name',
                                     custom_getter=_self._custom_dtype_getter)
    
     def forward(self, inputs, ...):
         with self._model_variable_scope():
             ...
             return outputs

    ```

  3. 放大loss,缩小gradient 过程比较简单,在使用optimizer计算梯度的时候将loss乘上loss_scale, 梯度计算完成后再将各个Variables的梯度缩小相应的倍数即可。 需要注意的地方在于loss计算过程最好是在FP32精度下进行,否则某些中间结果可能会超出FP16精度的范围, 如计算l2 loss的时候,对Variable值进行累加时和可能会上溢。

    def model_function(features, labels, mode, params):
        model = Model(dtype=tf.float16)
        logits = model.forward(features)
    
        # 结果转为32位再计算loss, 否则可能上溢
        cross_entropy = tf.losses.softmax_cross_entropy(labels,
                                                        tf.cast(logits, tf.float32))
        l2_loss = tf.nn.l2_loss(tf.cast(trainable_variables, tf.float32))
        total_loss = cross_entropy + l2_weight * l2_loss
    
        if mode == tf.estimator.ModeKeys.TRAIN:
            global_step = tf.train.get_or_create_globa_step()
            optimizer = tf.train.SomeOptimizer(...)
            scaled_grad_vars = optimizer.compute_gradients(total_loss * loss_scale) 
            unscaled_grad_vars = [(grad  / loss_scale, var) in scaled_grad_vars
                                    for grad, var in scaled_grad_vars]
            train_op = optimizer.apply_gradients(unscaled_grad_vars, global_step)
    
            return train_op

PyTorch

NVIDIA为PyTorch编写了一个用于混合精度训练的库apex, 安装好该库之后可以很方便的使用混合精度加速训练。

下面的代码主要参考https://github.com/NVIDIA/apex/tree/master/examples/imagenet

  1. 将输入转为FP16精度

    inputs = inputs.half()    
  2. 将模型转为FP16精度

    对于PyTorch中继承自torch.nn.Module的模型,直接调用器half()函数即可将其转为FP16精度。 但是需要注意的是,BN层需在FP32精度下进行计算。 若模型中包含BN层,需将其转回FP32。

    def BN_convert_float(module):
        if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
            module.float()
        for child in module.children():
            BN_convert_float(child)
        return module
    
    model = build_model()
    model = BN_convert_float(model.half())
  3. 使用FP16_Optimizer装饰优化器

    apex库提供的FP16_Optimizer可以帮助用户实现FP32 master copy功能。

    from apex.fp16_utils import FP16_Optimizer
    
    optimizer = torch.optim.SomeOptimizer(model.parameters(), ...)
    optimizer = FP16_Optimizer(optimizer,
                                static_loss_scale=loss_scale)

    在优化过程中,将loss.backward()替换为optimizer.backward(loss)即可实现loss放大与梯度缩小。

    while training:
        loss = ...
    
        optimizer.zero_grad()
        if fp16_mode:
            optimizer.backward(loss)
        else:
            loss.backward()
        optimizer.step()

参考

  1. Mixed Precision Training
  2. NVIDIA deep learning SDK documentation
  3. TensorFlow benchmarks
  4. TensorFlow models
  5. apex