备注
点击:ref:`此处<sphx_glr_download_recipes_recipes_amp_recipe.py>`下载完整示例代码
自动混合精度¶
Created On: Sep 15, 2020 | Last Updated: Jan 30, 2025 | Last Verified: Nov 05, 2024
作者: Michael Carilli
torch.cuda.amp 提供了便捷方法用于混合精度，其中某些操作使用``torch.float32``（float）数据类型，而其他操作使用``torch.float16``（half）数据类型。一些操作，例如线性层和卷积层，在``float16``或``bfloat16``中运行速度更快。而其他操作，例如归约操作，通常需要``float32``的动态范围。混合精度尝试将每个操作匹配到其适当的数据类型，这可以减少网络的运行时间和内存占用。
通常，“自动混合精度训练”一起使用`torch.autocast <https://pytorch.org/docs/stable/amp.html#torch.autocast>`_和`torch.cuda.amp.GradScaler <https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler>`_。
此配方测量了简单网络在默认精度下的性能，然后逐步添加``autocast``和``GradScaler``以在改进性能的混合精度下运行相同的网络。
您可以下载并将此配方作为独立的Python脚本运行。唯一的要求是PyTorch 1.6或更高版本以及支持CUDA的GPU。
混合精度主要受益于启用了Tensor Core的架构（Volta、Turing、Ampere）。此配方在这些架构上应显示显著（2-3倍）的速度提升。在较早的架构（Kepler、Maxwell、Pascal）上，您可能观察到温和的速度提升。运行``nvidia-smi``以显示您的GPU架构。
import torch, time, gc
# Timing utilities
start_time = None
def start_timer():
global start_time
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.synchronize()
start_time = time.time()
def end_timer_and_print(local_msg):
torch.cuda.synchronize()
end_time = time.time()
print("\n" + local_msg)
print("Total execution time = {:.3f} sec".format(end_time - start_time))
print("Max memory used by tensors = {} bytes".format(torch.cuda.max_memory_allocated()))
一个简单的网络¶
以下线性层和ReLU序列在混合精度下应显示速度提升。
def make_model(in_size, out_size, num_layers):
layers = []
for _ in range(num_layers - 1):
layers.append(torch.nn.Linear(in_size, in_size))
layers.append(torch.nn.ReLU())
layers.append(torch.nn.Linear(in_size, out_size))
return torch.nn.Sequential(*tuple(layers)).cuda()
batch_size、in_size、``out_size``和``num_layers``的选择足够大以让GPU饱和工作。通常，混合精度在GPU饱和时提供最大的速度提升。小型网络可能受CPU约束，在这种情况下混合精度无法提高性能。尺寸还选择了线性层的参与维度是8的倍数，以允许Tensor Core在支持Tensor Core的GPU上使用（请参阅:ref:故障排除<troubleshooting>）。
练习：改变参与尺寸，看看混合精度速度提升的变化。
batch_size = 512 # Try, for example, 128, 256, 513.
in_size = 4096
out_size = 4096
num_layers = 3
num_batches = 50
epochs = 3
device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.set_default_device(device)
# Creates data in default precision.
# The same data is used for both default and mixed precision trials below.
# You don't need to manually change inputs' ``dtype`` when enabling mixed precision.
data = [torch.randn(batch_size, in_size) for _ in range(num_batches)]
targets = [torch.randn(batch_size, out_size) for _ in range(num_batches)]
loss_fn = torch.nn.MSELoss().cuda()
默认精度¶
没有``torch.cuda.amp``时，以下简单网络在默认精度（torch.float32）下执行所有操作：
net = make_model(in_size, out_size, num_layers)
opt = torch.optim.SGD(net.parameters(), lr=0.001)
start_timer()
for epoch in range(epochs):
for input, target in zip(data, targets):
output = net(input)
loss = loss_fn(output, target)
loss.backward()
opt.step()
opt.zero_grad() # set_to_none=True here can modestly improve performance
end_timer_and_print("Default precision:")
添加``torch.autocast``¶
实例`torch.autocast <https://pytorch.org/docs/stable/amp.html#autocasting>`_作为上下文管理器，允许脚本区域以混合精度运行。
在这些区域中，CUDA操作在由``autocast``选择的``dtype``中运行，以提高性能并保持准确性。有关每个操作选择的精度详情以及在何种情况下，请参阅`Autocast操作参考 <https://pytorch.org/docs/stable/amp.html#autocast-op-reference>`_。
for epoch in range(0): # 0 epochs, this section is for illustration only
for input, target in zip(data, targets):
# Runs the forward pass under ``autocast``.
with torch.autocast(device_type=device, dtype=torch.float16):
output = net(input)
# output is float16 because linear layers ``autocast`` to float16.
assert output.dtype is torch.float16
loss = loss_fn(output, target)
# loss is float32 because ``mse_loss`` layers ``autocast`` to float32.
assert loss.dtype is torch.float32
# Exits ``autocast`` before backward().
# Backward passes under ``autocast`` are not recommended.
# Backward ops run in the same ``dtype`` ``autocast`` chose for corresponding forward ops.
loss.backward()
opt.step()
opt.zero_grad() # set_to_none=True here can modestly improve performance
添加``GradScaler``¶
`梯度缩放 <https://pytorch.org/docs/stable/amp.html#gradient-scaling>`_帮助防止梯度在混合精度训练中因幅度小而归零（“下溢”）。
`torch.cuda.amp.GradScaler <https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler>`_方便地执行梯度缩放步骤。
# Constructs a ``scaler`` once, at the beginning of the convergence run, using default arguments.
# If your network fails to converge with default ``GradScaler`` arguments, please file an issue.
# The same ``GradScaler`` instance should be used for the entire convergence run.
# If you perform multiple convergence runs in the same script, each run should use
# a dedicated fresh ``GradScaler`` instance. ``GradScaler`` instances are lightweight.
scaler = torch.amp.GradScaler("cuda")
for epoch in range(0): # 0 epochs, this section is for illustration only
for input, target in zip(data, targets):
with torch.autocast(device_type=device, dtype=torch.float16):
output = net(input)
loss = loss_fn(output, target)
# Scales loss. Calls ``backward()`` on scaled loss to create scaled gradients.
scaler.scale(loss).backward()
# ``scaler.step()`` first unscales the gradients of the optimizer's assigned parameters.
# If these gradients do not contain ``inf``s or ``NaN``s, optimizer.step() is then called,
# otherwise, optimizer.step() is skipped.
scaler.step(opt)
# Updates the scale for next iteration.
scaler.update()
opt.zero_grad() # set_to_none=True here can modestly improve performance
汇总：“自动混合精度”¶
(以下还演示了``enabled``，这是``autocast``和``GradScaler``的一个可选便捷参数。如果为False，``autocast``和``GradScaler``的调用将变为空操作。这允许在默认精度和混合精度之间切换，而不需要if/else语句。)
use_amp = True
net = make_model(in_size, out_size, num_layers)
opt = torch.optim.SGD(net.parameters(), lr=0.001)
scaler = torch.amp.GradScaler("cuda" ,enabled=use_amp)
start_timer()
for epoch in range(epochs):
for input, target in zip(data, targets):
with torch.autocast(device_type=device, dtype=torch.float16, enabled=use_amp):
output = net(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
scaler.step(opt)
scaler.update()
opt.zero_grad() # set_to_none=True here can modestly improve performance
end_timer_and_print("Mixed precision:")
检查/修改梯度（例如裁剪）¶
所有由``scaler.scale(loss).backward()``产生的梯度都被缩放。如果您希望在``backward()``和``scaler.step(optimizer)``之间修改或检查参数的``.grad``属性，您应该先使用`scaler.unscale_(optimizer) <https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler.unscale_>`_取消缩放它们。
for epoch in range(0): # 0 epochs, this section is for illustration only
for input, target in zip(data, targets):
with torch.autocast(device_type=device, dtype=torch.float16):
output = net(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
# Unscales the gradients of optimizer's assigned parameters in-place
scaler.unscale_(opt)
# Since the gradients of optimizer's assigned parameters are now unscaled, clips as usual.
# You may use the same value for max_norm here as you would without gradient scaling.
torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=0.1)
scaler.step(opt)
scaler.update()
opt.zero_grad() # set_to_none=True here can modestly improve performance
保存/恢复¶
要以位级精度保存/恢复支持Amp的运行，请使用`scaler.state_dict <https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler.state_dict>`_和`scaler.load_state_dict <https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler.load_state_dict>`_。
保存时，将``scaler``状态字典与通常的模型和优化器状态``dicts``一起保存。在迭代开始时的任意位置或在``scaler.update()``之后的迭代结束时执行此操作。
checkpoint = {"model": net.state_dict(),
"optimizer": opt.state_dict(),
"scaler": scaler.state_dict()}
# Write checkpoint as desired, e.g.,
# torch.save(checkpoint, "filename")
恢复时，将``scaler``状态字典与模型和优化器状态``dicts``一起加载。按需读取检查点，例如：
dev = torch.cuda.current_device()
checkpoint = torch.load("filename",
map_location = lambda storage, loc: storage.cuda(dev))
net.load_state_dict(checkpoint["model"])
opt.load_state_dict(checkpoint["optimizer"])
scaler.load_state_dict(checkpoint["scaler"])
如果一个检查点是从不带Amp的运行中创建的，并且您希望在带Amp的情况下恢复训练，请像往常一样从检查点加载模型和优化器状态。检查点不会包含保存的``scaler``状态，因此请使用一个新的``GradScaler``实例。
如果一个检查点是从带Amp的运行中创建的，而您希望在不带``Amp``的情况下恢复训练，请像往常一样从检查点加载模型和优化器状态，并忽略保存的``scaler``状态。
推理/评估¶
autocast``可以单独用于包装推理或评估的前向过程。不需要``GradScaler。
高级主题¶
有关高级用例，请参阅`自动混合精度示例 <https://pytorch.org/docs/stable/notes/amp_examples.html>`_，包括：
梯度累积
梯度惩罚/双向传播
包含多个模型、优化器或损失的网络
多GPU（
torch.nn.DataParallel``或``torch.nn.parallel.DistributedDataParallel）自定义自动求导函数（``torch.autograd.Function``的子类）
如果您在同一脚本中执行多个收敛运行，每次运行应使用一个专门的新的``GradScaler``实例。``GradScaler``实例是轻量级的。
如果您正在使用调度器注册自定义C++操作，请参阅调度器教程的`autocast部分 <https://pytorch.org/tutorials/advanced/dispatcher.html#autocast>`_。
故障排除¶
使用Amp的速度提升较小¶
您的网络可能未能通过工作量饱和GPU，因此受到CPU约束。Amp对GPU性能的影响将不会显现。
一个粗略的经验法则是尽可能增加批量大小和/或网络尺寸，而不会发生OOM。
尽量避免过多的CPU-GPU同步（例如``.item()``调用或从CUDA张量打印值）。
尽量避免许多小型CUDA操作的序列（如果可以，将它们合并为少量大型CUDA操作）。
您的网络可能是GPU计算约束型（具有大量``matmuls``/卷积操作），但您的GPU没有Tensor Cores。在这种情况下，预期速度提升会减少。
``matmul``维度不适合Tensor Core。确保``matmuls``的参与尺寸是8的倍数。（对于带有编码器/解码器的NLP模型，这可能会很微妙。此外，卷积曾经对Tensor Core使用有类似的尺寸约束，但对于CuDNN版本7.3及以上，没有这样的约束。参见`此处 <https://github.com/NVIDIA/apex/issues/221#issuecomment-478084841>`_以获取指导。）
损失是inf/NaN¶
首先，检查您的网络是否符合:ref:高级用例<advanced-topics>。另见`推荐使用binary_cross_entropy_with_logits而不是binary_cross_entropy <https://pytorch.org/docs/stable/amp.html#prefer-binary-cross-entropy-with-logits-over-binary-cross-entropy>`_。
如果您确信您的Amp使用是正确的，可能需要提交一个问题，但在提交之前，收集以下信息会很有帮助：
分别禁用``autocast``或``GradScaler``（通过将``enabled=False``传递给它们的构造函数），看看``infs``/``NaNs``是否仍然存在。
如果您怀疑网络的一部分（例如，一个复杂的损失函数）溢出，请使用``float32``运行该区域的前向传播，看看是否仍然存在``inf``或``NaN``。`autocast文档字符串 <https://pytorch.org/docs/stable/amp.html#torch.autocast>`_中的最后一个代码片段显示了如何强制一个子区域以``float32``运行（通过局部禁用``autocast``并对该子区域的输入进行类型转换）。
