提示:文章写完后,目录可以自动生成,如何生成可参考右边的帮助文档
前言
提示:这里可以添加本文要记录的大概内容:
例如:随着人工智能的不断发展,机器学习这门技术也越来越重要,很多人都开启了学习机器学习,本文就介绍了机器学习的基础内容。
提示:以下是本篇文章正文内容,下面案例可供参考
一、混合精度训练
原理:
将FP32的数据存储成FP16显然可以减小存储空间也可加快运算速度。但对于 FP16 不可避免的问题就是:表示的范围较窄,如下图所示,大量非 0 梯度会遇到溢出问题。
解决办法是:
对梯度乘一个 [公式] 的系数,称为 scale factor,把梯度 shift 到 FP16 的表示范围。而这一切,自动混合精度训练都给我们安排好了。
步骤:
GradScaler的工作就是在反向传播前给 loss 乘一个 scale factor,所以之后反向传播得到的梯度都乘了相同的 scale factor。并且为了不影响学习率,在梯度更新前将梯度unscale。总结amp的基本训练流程:
- 维护一个 FP32 数值精度模型的副本
- 在每个iteration
1.拷贝并且转换成 FP16 模型
2.前向传播(FP16 的模型参数)
3.loss 乘 scale factor s
4.反向传播(FP16 的模型参数和参数梯度)
5.参数梯度乘 1/s
6.利用 FP16 的梯度更新 FP32 的模型参数
二、使用步骤
1.引入库
代码如下(示例):
import torch
2.使用混合精度训练
代码如下(示例):
# amp依赖Tensor core架构,所以model参数必须是cuda tensor类型
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# GradScaler对象用来自动做梯度缩放
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
# 在autocast enable 区域运行forward
with autocast():
# model做一个FP16的副本,forward
output = model(input)
loss = loss_fn(output, target)
# 用scaler,scale loss(FP16),backward得到scaled的梯度(FP16)
scaler.scale(loss).backward()
# scaler 更新参数,会先自动unscale梯度
# 如果有nan或inf,自动跳过
scaler.step(optimizer)
# scaler factor更新
scaler.update()
总结
本文仅仅简单介绍了混合精度训练的使用,混合精度训练能够在现有硬件的情况下,为我们节省内存并提升模型的运算速度。