昇思MindSpore学习入门-自动混合精度

混合精度（Mix Precision）训练是指在训练时，对神经网络不同的运算采用不同的数值精度的运算策略。在神经网络运算中，部分运算对数值精度不敏感，此时使用较低精度可以达到明显的加速效果（如conv、matmul等）；而部分运算由于输入和输出的数值差异大，通常需要保留较高精度以保证结果的正确性（如log、softmax等）。

当前的AI加速卡通常通过针对计算密集、精度不敏感的运算设计了硬件加速模块，如NVIDIA GPU的TensorCore、Ascend NPU的Cube等。对于conv、matmul等运算占比较大的神经网络，其训练速度通常会有较大的加速比。

mindspore.amp模块提供了便捷的自动混合精度接口，用户可以在不同的硬件后端通过简单的接口调用获得训练加速。下面我们对混合精度计算原理进行简介，而后通过实例介绍MindSpore的自动混合精度用法。

混合精度计算原理

浮点数据类型主要分为双精度（FP64）、单精度（FP32）、半精度（FP16）。在神经网络模型的训练过程中，一般默认采用单精度（FP32）浮点数据类型，来表示网络模型权重和其他参数。在了解混合精度训练之前，这里简单了解浮点数据类型。

根据IEEE二进制浮点数算术标准（IEEE 754）的定义，浮点数据类型分为双精度（FP64）、单精度（FP32）、半精度（FP16）三种，其中每一种都有三个不同的位来表示。FP64表示采用8个字节共64位，来进行的编码存储的一种数据类型；同理，FP32表示采用4个字节共32位来表示；FP16则是采用2字节共16位来表示。如图所示：

从图中可以看出，与FP32相比，FP16的存储空间是FP32的一半。类似地，FP32则是FP64的一半。因此使用FP16进行运算具备以下优势：

• 减少内存占用：FP16的位宽是FP32的一半，因此权重等参数所占用的内存也是原来的一半，节省下来的内存可以放更大的网络模型或者使用更多的数据进行训练。
• 计算效率更高：在特殊的AI加速芯片如华为Atlas训练系列产品和Atlas 200/300/500推理产品系列，或者NVIDIA VOLTA架构的GPU上，使用FP16的执行运算性能比FP32更加快。
• 加快通讯效率：针对分布式训练，特别是在大模型训练的过程中，通讯的开销制约了网络模型训练的整体性能，通讯的位宽少了意味着可以提升通讯性能，减少等待时间，加快数据的流通。

但是使用FP16同样会带来一些问题：

• 数据溢出：FP16的有效数据表示范围为 [5.9×10−8,65504]，FP32的有效数据表示范围为 [1.4×10−45,1.7×1038]。可见FP16相比FP32的有效范围要窄很多，使用FP16替换FP32会出现上溢（Overflow）和下溢（Underflow）的情况。而在深度学习中，需要计算网络模型中权重的梯度（一阶导数），因此梯度会比权重值更加小，往往容易出现下溢情况。
• 舍入误差：Rounding Error是指当网络模型的反向梯度很小，一般FP32能够表示，但是转换到FP16会小于当前区间内的最小间隔，会导致数据溢出。如0.00006666666在FP32中能正常表示，转换到FP16后会表示成为0.000067，不满足FP16最小间隔的数会强制舍入。
• 因此，在使用混合精度获得训练加速和内存节省的同时，需要考虑FP16引入问题的解决。Loss Scale损失缩放，FP16类型数据下溢问题的解决方案，其主要思想是在计算损失值loss的时候，将loss扩大一定的倍数。根据链式法则，梯度也会相应扩大，然后在优化器更新权重时再缩小相应的倍数，从而避免了数据下溢。
• 根据上述原理介绍，典型的混合精度计算流程如下图所示：

1. 参数以FP32存储；
2. 正向计算过程中，遇到FP16算子，需要把算子输入和参数从FP32 cast成FP16进行计算；
3. 将Loss层设置为FP32进行计算；
4. 反向计算过程中，首先乘以Loss Scale值，避免反向梯度过小而产生下溢；
5. FP16参数参与梯度计算，其结果将被cast回FP32；
6. 除以Loss scale值，还原被放大的梯度；
7. 判断梯度是否存在溢出，如果溢出则跳过更新，否则优化器以FP32对原始参数进行更新。

下面我们通过导入快速入门中的手写数字识别模型及数据集，演示MindSpore的自动混合精度实现。

类型转换

混合精度计算需要将需要使用低精度的运算进行类型转换，将其输入转为FP16类型，得到输出后进将其重新转回FP32类型。MindSpore同时提供了自动和手动类型转换的方法，满足对易用性和灵活性的不同需求，下面我们分别对其进行介绍。

自动类型转换

mindspore.amp.auto_mixed_precision 接口提供对网络做自动类型转换的功能。自动类型转换遵循黑白名单机制，根据常用的运算精度习惯配置了4个等级，分别为：

• ‘O0’：神经网络保持FP32；
• ‘O1’：按白名单将运算cast为FP16；
• ‘O2’：按黑名单保留FP32，其余运算cast为FP16；
• ‘O3’：神经网络完全cast为FP16。

下面是使用自动类型转换的示例：

from mindspore.amp import auto_mixed_precision

model = Network()

model = auto_mixed_precision(model, 'O2')

手动类型转换

通常情况下自动类型转换可以通过满足大部分混合精度训练的需求，但当用户需要精细化控制神经网络不同部分的运算精度时，可以通过手动类型转换的方式进行控制。

Cell粒度类型转换

nn.Cell类提供了to_float方法，可以一键配置该模块的运算精度，自动将模块输入cast为指定的精度：

class NetworkFP16(nn.Cell):

    def __init__(self):

        super().__init__()

        self.flatten = nn.Flatten()

        self.dense_relu_sequential = nn.SequentialCell(

            nn.Dense(28*28, 512).to_float(ms.float16),

            nn.ReLU(),

            nn.Dense(512, 512).to_float(ms.float16),

            nn.ReLU(),

            nn.Dense(512, 10).to_float(ms.float16)

        )

    def construct(self, x):

        x = self.flatten(x)

        logits = self.dense_relu_sequential(x)

        return logits

自定义粒度类型转换

当用户需要在单个运算，或多个模块组合配置运算精度时，Cell粒度往往无法满足，此时可以直接通过对输入数据的类型进行cast来达到自定义粒度控制的目的。

class NetworkFP16Manual(nn.Cell):

    def __init__(self):

        super().__init__()

        self.flatten = nn.Flatten()

        self.dense_relu_sequential = nn.SequentialCell(

            nn.Dense(28*28, 512),

            nn.ReLU(),

            nn.Dense(512, 512),

            nn.ReLU(),

            nn.Dense(512, 10)

        )

    def construct(self, x):

        x = self.flatten(x)

        x = x.astype(ms.float16)

        logits = self.dense_relu_sequential(x)

        logits = logits.astype(ms.float32)

        return logits

损失缩放

MindSpore中提供了两种Loss Scale的实现，分别为StaticLossScaler和DynamicLossScaler，其差异为损失缩放值scale value是否进行动态调整。下面以DynamicLossScalar为例，根据混合精度计算流程实现神经网络训练逻辑。

首先，实例化LossScaler，并在定义前向网络时，手动放大loss值。

from mindspore.amp import DynamicLossScaler

# Instantiate loss function and optimizer

loss_fn = nn.CrossEntropyLoss()

optimizer = nn.SGD(model.trainable_params(), 1e-2)

# Define LossScaler

loss_scaler = DynamicLossScaler(scale_value=2**16, scale_factor=2, scale_window=50)

def forward_fn(data, label):

    logits = model(data)

    loss = loss_fn(logits, label)

    # scale up the loss value

    loss = loss_scaler.scale(loss)

return loss, logits

接下来进行函数变换，获得梯度函数。

grad_fn = value_and_grad(forward_fn, None, model.trainable_params())

定义训练step：计算当前梯度值并恢复损失。使用 all_finite 判断是否出现梯度下溢问题，如果无溢出，恢复梯度并更新网络权重；如果溢出，跳过此step。

from mindspore.amp import all_finite

@ms.jit

def train_step(data, label):

    (loss, _), grads = grad_fn(data, label)

    loss = loss_scaler.unscale(loss)

    is_finite = all_finite(grads)

    if is_finite:

        grads = loss_scaler.unscale(grads)

        optimizer(grads)

    loss_scaler.adjust(is_finite)

    return loss

最后，我们训练1个epoch，观察使用自动混合精度训练的loss收敛情况。

size = train_dataset.get_dataset_size()

model.set_train()

for batch, (data, label) in enumerate(train_dataset.create_tuple_iterator()):

    loss = train_step(data, label)

    if batch % 100 == 0:

        loss, current = loss.asnumpy(), batch

        print(f"loss: {loss:>7f}  [{current:>3d}/{size:>3d}]")

可以看到loss收敛趋势正常，没有出现溢出问题。

Cell配置自动混合精度

MindSpore支持使用Cell封装完整计算图的编程范式，此时可以使用mindspore.amp.build_train_network接口，自动进行类型转换，并将Loss Scale传入，作为整图计算的一部分。 此时仅需要配置混合精度等级和LossScaleManager即可获得配置好自动混合精度的计算图。

FixedLossScaleManager和DynamicLossScaleManager是Cell配置自动混合精度的Loss scale管理接口，分别与StaticLossScalar和DynamicLossScalar对应，具体详见mindspore.amp。

from mindspore.amp import build_train_network, FixedLossScaleManager

model = Network()

loss_scale_manager = FixedLossScaleManager()

model = build_train_network(model, optimizer, loss_fn, level="O2", loss_scale_manager=loss_scale_manager)

Model 配置自动混合精度

mindspore.train.Model是神经网络快速训练的高阶封装，其将mindspore.amp.build_train_network封装在内，因此同样只需要配置混合精度等级和LossScaleManager，即可进行自动混合精度训练。