Learn AI with Python · 课时

使用 AMP 进行混合精度训练

torch.cuda.amp.autocast()、GradScaler、FP16 与 BF16、内存节省和速度提升。

使用 AMP 进行混合精度训练 是 CoddyKit 上的免费 Learn AI with Python 课时。这是第 3 节课，共 4 节。你可以在下方免费阅读本课时的完整内容 — 然后在浏览器中使用内置代码编辑器和全天候 AI 导师进行实践。这是 Learn AI with Python 学习路径的一部分，你的进度在网页和 CoddyKit 应用中同步。 Learn AI with Python 课程共包含 4 节课。

什么是混合精度

混合精度训练会使用 16 位浮点数（FP16）而不是 32 位浮点数（FP32）运行大多数操作。FP16 只需一半的内存，并且可以在现代 GPU 张量核心上运行得更快；同时，少数敏感操作会保留 FP32，以确保稳定性。

优势

混合精度可以带来以下优势：

节省约 2 倍内存，从而支持更大的批次或模型
在张量核心上进行更快的矩阵乘法
正确使用时，最终模型精度几乎不会下降或完全不会下降

FP16 下溢问题

FP16 的数值范围较窄。较小的梯度值可能会下溢为零，从而在不易察觉的情况下停止学习。这是混合精度必须解决的核心挑战，也是我们不能简单地将所有内容都转换为 FP16 的原因。

torch.cuda.amp

PyTorch 的自动混合精度（AMP）通过两个工具处理这些细节：autocast会为每个操作选择精度，而 GradScaler则防止梯度下溢。

import torch
from torch.cuda.amp import autocast, GradScaler

autocast 上下文

请将前向传播包装在 autocast() 中。在其中，PyTorch 会自动以最安全的精度运行每个操作：矩阵乘法使用 FP16，而 softmax 和损失等归约操作使用 FP32。

with autocast():
    output = model(x)
    loss = criterion(output, y)

引入 GradScaler

GradScaler会在反向传播之前将损失乘以一个较大的 scale 因子，从而防止下溢。这样可以将微小梯度移动到 FP16 可表示的范围内；在优化器 step 之前会移除该 scale。

scaler = GradScaler()

缩放损失

scaler.scale(loss).backward()会放大损失，然后执行反向传播。生成的梯度也会经过缩放，从而避免较小的值消失。

scaler.scale(loss).backward()

scaler.step

scaler.step(optimizer)会先将梯度反缩放为真实幅值，然后调用 optimizer.step()，但前提是没有出现无穷值或 NaNs。如果梯度发生溢出，则会跳过这次 step。

scaler.step(optimizer)

scaler.update

scaler.update()会调整下一次迭代使用的 scale 因子：当步骤成功时增大 scale，发生溢出后减小 scale。这种自适应缩放会自动保持训练稳定。

scaler.update()

完整的 AMP 循环

将这些部分组合起来，就得到一个完整的混合精度训练步骤。

scaler = GradScaler()
for x, y in loader:
    x, y = x.cuda(), y.cuda()
    optimizer.zero_grad()
    with autocast():
        out = model(x)
        loss = criterion(out, y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

实用提示

请牢记以下几点：

只对经过缩放的损失调用 backward()
autocast 只包装前向传播，不包装 backward 或优化器步骤
AMP 在采用张量核心的 GPU 上效果显著；在较旧的硬件上，收益会更小

快速检查

请测试您对 AMP 的掌握情况。

回顾

您已经学习了使用 AMP 进行混合精度训练：

autocast()会在前向传播中为每个操作选择 FP16 或 FP32
GradScaler会缩放损失，以避免梯度下溢
完整流程是 scaler.scale(loss).backward()、scaler.step(optimizer)、scaler.update()
结果：大约节省 2 倍内存，并加快训练

常见问题解答

「使用 AMP 进行混合精度训练」课时是免费的吗？

是的 — 「使用 AMP 进行混合精度训练」的完整文本可在网页上免费阅读。要进行交互式练习（内置代码编辑器和全天候 AI 导师）并解锁 Learn AI with Python 课程的其余内容，请升级到 CoddyKit PRO。 Learn AI with Python 课程共包含 4 节课。

「使用 AMP 进行混合精度训练」这节课中我会学到什么？

torch.cuda.amp.autocast()、GradScaler、FP16 与 BF16、内存节省和速度提升。你通过在浏览器中直接运行的动手代码来练习 Learn AI with Python，全天候 AI 导师会在你学习这节课的过程中回答你的问题。

学习 Learn AI with Python 需要有经验吗？

无需任何先前经验。CoddyKit 上的 Learn AI with Python 课程适合初学者到高级学习者，你可以从这里开始或从头开始，按照自己的节奏学习。这是第 3 节课，共 4 节。

「使用 AMP 进行混合精度训练」课时需要多长时间？

大多数 CoddyKit 课程大约需要 5–10 分钟。每节课都很精短且互动，所以你能稳步进步，并在网页和应用中从离开的地方继续。

我能在这节 Learn AI with Python 课中编写并运行代码吗？

能。每节 Learn AI with Python 课都包含内置代码编辑器，你可以在浏览器中直接编写并运行真实代码，并获得即时 AI 反馈 — 无需本地设置。

此课程中的所有课时

← 返回 Learn AI with Python