大家好，我是刘明，明志科技创始人，华为昇思MindSpore布道师。
技术上主攻前端开发、鸿蒙开发和AI算法研究。
努力为大家带来持续的技术分享，如果你也喜欢我的文章，就点个关注吧
正文开始
混合精度（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引入问题的解决。Loss Scale损失缩放，FP16类型数据下溢问题的解决方案，其主要思想是在计算损失值loss的时候，将loss扩大一定的倍数。根据链式法则，梯度也会相应扩大，然后在优化器更新权重时再缩小相应的倍数，从而避免了数据下溢。
根据上述原理介绍，典型的混合精度计算流程如下图所示：
- 参数以FP32存储；
- 正向计算过程中，遇到FP16算子，需要把算子输入和参数从FP32 cast成FP16进行计算；
- 将Loss层设置为FP32进行计算；
- 反向计算过程中，首先乘以Loss Scale值，避免反向梯度过小而产生下溢；
- FP16参数参与梯度计算，其结果将被cast回FP32；
- 除以Loss scale值，还原被放大的梯度；
- 判断梯度是否存在溢出，如果溢出则跳过更新，否则优化器以FP32对原始参数进行更新。
下面我们通过导入快速入门中的手写数字识别模型及数据集，演示MindSpore的自动混合精度实现。
import mindspore as ms
from mindspore import nn
from mindspore import value_and_grad
from mindspore.dataset import vision, transforms
from mindspore.dataset import MnistDataset
# Download data from open datasets
from download import download
url = "https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/" \
"notebook/datasets/MNIST_Data.zip"
path = download(url, "./", kind="zip", replace=True)
def datapipe(path, batch_size):
image_transforms = [
vision.Rescale(1.0 / 255.0, 0),
vision.Normalize(mean=(0.1307,), std=(0.3081,)),
vision.HWC2CHW()
]
label_transform = transforms.TypeCast(ms.int32)
dataset = MnistDataset(path)
dataset = dataset.map(image_transforms, 'image')
dataset = dataset.map(label_transform, 'label')
dataset = dataset.batch(batch_size)
return dataset
train_dataset = datapipe('MNIST_Data/train', 64)
# Define model
class Network(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)
logits = self.dense_relu_sequential(x)
return logits
类型转换
混合精度计算需要将需要使用低精度的运算进行类型转换，将其输入转为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')
手动类型转换
通常情况下自动类型转换可以通过满足大部分混合精度训练的需求，但当用户需要精细化控制神经网络不同部分的运算精度时，可以通过手动类型转换的方式进行控制。
下面我们对前文的Network进行改造，演示手动类型转换的不同方式。
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(), has_aux=True)
定义训练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，即可进行自动混合精度训练。
from mindspore.train import Model, LossMonitor
# Initialize network
model = Network()
loss_fn = nn.CrossEntropyLoss()
optimizer = nn.SGD(model.trainable_params(), 1e-2)
loss_scale_manager = FixedLossScaleManager()
trainer = Model(model, loss_fn=loss_fn, optimizer=optimizer, metrics={'accuracy'}, amp_level="O2", loss_scale_manager=loss_scale_manager)
loss_callback = LossMonitor(100)
trainer.train(10, train_dataset, callbacks=[loss_callback])
扫一扫
313
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
