淘先锋技术网

首页 1 2 3 4 5 6 7

提示:文章写完后,目录可以自动生成,如何生成可参考右边的帮助文档


前言

提示:这里可以添加本文要记录的大概内容:

例如:随着人工智能的不断发展,机器学习这门技术也越来越重要,很多人都开启了学习机器学习,本文就介绍了机器学习的基础内容。


提示:以下是本篇文章正文内容,下面案例可供参考

一、混合精度训练

原理:

将FP32的数据存储成FP16显然可以减小存储空间也可加快运算速度。但对于 FP16 不可避免的问题就是:表示的范围较窄,如下图所示,大量非 0 梯度会遇到溢出问题。

解决办法是:

对梯度乘一个 [公式] 的系数,称为 scale factor,把梯度 shift 到 FP16 的表示范围。而这一切,自动混合精度训练都给我们安排好了。

步骤:

GradScaler的工作就是在反向传播前给 loss 乘一个 scale factor,所以之后反向传播得到的梯度都乘了相同的 scale factor。并且为了不影响学习率,在梯度更新前将梯度unscale。总结amp的基本训练流程:

  1. 维护一个 FP32 数值精度模型的副本
  2. 在每个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()

总结

本文仅仅简单介绍了混合精度训练的使用,混合精度训练能够在现有硬件的情况下,为我们节省内存并提升模型的运算速度。