CUDA 自动混合精度示例¶
通常，“自动混合精度训练”意味着使用
torch.autocast 和 torch.cuda.amp.GradScaler 一起进行训练。
实例 torch.autocast 为选定区域启用自动转换。
自动转换会自动选择 GPU 操作的精度，以提高性能
同时保持准确性。
实例 torch.cuda.amp.GradScaler 有助于方便地执行梯度缩放的步骤。梯度缩放通过最小化梯度下溢来改善具有 float16 梯度的网络的收敛性，如此处所述。
torch.autocast 和 torch.cuda.amp.GradScaler 是模块化的。
在下面的示例中，每个都按照其单独的文档建议使用。
（这里的示例仅供参考。请参阅 自动混合精度配方 以获取可运行的演练。）
典型的混合精度训练¶
# 在默认精度下创建模型和优化器
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 在训练开始时创建一个GradScaler。
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
# 使用自动混合精度运行前向传播。
with autocast(device_type='cuda', dtype=torch.float16):
output = model(input)
loss = loss_fn(output, target)
# 缩放损失。在缩放后的损失上调用backward()以创建缩放后的梯度。
# 不推荐在自动混合精度下进行反向传播。
# 反向操作运行在与前向操作对应的自动混合精度选择的相同数据类型下。
scaler.scale(loss).backward()
# scaler.step()首先对优化器分配的参数的梯度进行缩放。
# 如果这些梯度不包含inf或NaN，则调用optimizer.step()，
# 否则，跳过optimizer.step()。
scaler.step(optimizer)
# 更新下一次迭代的缩放比例。
scaler.update()
使用未缩放的梯度¶
由 scaler.scale(loss).backward() 生成的所有梯度都会被缩放。如果您希望在 backward() 和 scaler.step(optimizer) 之间修改或检查参数的 .grad 属性，您应该首先对它们进行反缩放。例如，梯度裁剪会操作一组梯度，使得它们的全局范数（参见 torch.nn.utils.clip_grad_norm_()）或最大幅度（参见 torch.nn.utils.clip_grad_value_()）
小于等于某个用户设定的阈值。如果您尝试在不反缩放的情况下进行裁剪，梯度的范数/最大幅度也会被缩放，因此您请求的阈值（原本是针对未缩放梯度的阈值）将无效。
scaler.unscale_(optimizer) 取消由 optimizer 分配的参数所持有的梯度缩放。
如果你的模型或模型包含分配给另一个优化器的其他参数（例如 optimizer2），你可以单独调用 scaler.unscale_(optimizer2) 来取消这些参数的梯度缩放。
梯度裁剪¶
在裁剪之前调用 scaler.unscale_(optimizer) 可以使您像往常一样裁剪未缩放的梯度：
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
with autocast(device_type='cuda', dtype=torch.float16):
output = model(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
# 就地取消优化器分配参数的梯度缩放
scaler.unscale_(optimizer)
# 由于优化器分配参数的梯度已取消缩放，因此按常规方式裁剪：
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
# 优化器的梯度已经取消缩放，因此 scaler.step 不会再次取消缩放它们，
# 尽管它仍然会跳过 optimizer.step() 如果梯度包含 infs 或 NaNs。
scaler.step(optimizer)
# 更新下一次迭代的缩放比例。
scaler.update()
scaler 记录了在此迭代中已经为该优化器调用了 scaler.unscale_(optimizer)，因此 scaler.step(optimizer) 知道不要在调用 optimizer.step() 之前冗余地取消缩放梯度。
警告
unscale_ 每个优化器每次调用 step 时只能调用一次，
并且只能在为该优化器分配的所有参数的梯度累积之后调用。
在每次 step 之间为给定优化器调用 unscale_ 两次会触发 RuntimeError。
使用缩放梯度¶
梯度累积¶
梯度累积在大小为 batch_per_iter * iters_to_accumulate
（如果分布式则为 * num_procs）的有效批次上累积梯度。 比例应针对有效批次进行校准，这意味着需要进行inf/NaN检查，
如果在找到inf/NaN梯度时跳过步骤，并且比例更新应在有效批次的粒度上进行。
此外，梯度应保持缩放，并且比例因子应保持不变，而在累积给定有效批次的梯度时。 如果在累积完成之前梯度未缩放（或比例因子发生变化），
则在下一个反向传递中，将缩放后的梯度添加到未缩放的梯度（或按不同因子缩放的梯度）之后，无法恢复累积的未缩放梯度 step 必须应用。
因此，如果你想对梯度进行unscale_（例如，允许对未缩放的梯度进行裁剪），
请在所有即将进行的step的（缩放）梯度累积之后，在step之前调用unscale_。此外，仅在调用了step的完整有效批次的迭代结束时调用update：
scaler = GradScaler()
for epoch in epochs:
for i, (input, target) in enumerate(data):
with autocast(device_type='cuda', dtype=torch.float16):
output = model(input)
loss = loss_fn(output, target)
loss = loss / iters_to_accumulate
# 累加缩放后的梯度。
scaler.scale(loss).backward()
if (i + 1) % iters_to_accumulate == 0:
# 如果需要，可以在这里进行unscale操作（例如，允许剪裁未缩放的梯度）
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
梯度惩罚¶
梯度惩罚的实现通常使用
torch.autograd.grad() 创建梯度，将它们组合以生成惩罚值，
并将惩罚值添加到损失中。
这是一个没有梯度缩放或自动转换的L2惩罚的普通示例：
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
output = model(input)
loss = loss_fn(output, target)
# 创建梯度
grad_params = torch.autograd.grad(outputs=loss,
inputs=model.parameters(),
create_graph=True)
# 计算惩罚项并将其添加到损失中
grad_norm = 0
for grad in grad_params:
grad_norm += grad.pow(2).sum()
grad_norm = grad_norm.sqrt()
loss = loss + grad_norm
loss.backward()
# 如果需要，在此处裁剪梯度
optimizer.step()
要实现带有梯度缩放的梯度惩罚，传递给 torch.autograd.grad() 的 outputs 张量应进行缩放。因此，得到的梯度将被缩放，并且在组合以创建惩罚值之前应进行反缩放。
此外，惩罚项的计算是前向传播的一部分，因此应该在一个autocast上下文中进行。
以下是相同的L2惩罚的显示方式：
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
with autocast(device_type='cuda', dtype=torch.float16):
output = model(input)
loss = loss_fn(output, target)
# 为 autograd.grad 的反向传播缩放损失，生成 scaled_grad_params
scaled_grad_params = torch.autograd.grad(outputs=scaler.scale(loss),
inputs=model.parameters(),
create_graph=True)
# 在计算惩罚项之前创建未缩放的 grad_params。scaled_grad_params 不属于任何优化器，
# 因此使用普通除法而不是 scaler.unscale_：
inv_scale = 1./scaler.get_scale()
grad_params = [p * inv_scale for p in scaled_grad_params]
# 计算惩罚项并将其添加到损失中
with autocast(device_type='cuda', dtype=torch.float16):
grad_norm = 0
for grad in grad_params:
grad_norm += grad.pow(2).sum()
grad_norm = grad_norm.sqrt()
loss = loss + grad_norm
# 像往常一样对反向调用应用缩放。
# 累积正确缩放的叶子梯度。
scaler.scale(loss).backward()
# 如果需要，可以在这里进行 unscale_（例如，允许剪裁未缩放的梯度）
# step() 和 update() 像往常一样进行。
scaler.step(optimizer)
scaler.update()
使用多个模型、损失函数和优化器¶
如果你的网络有多个损失，你必须分别对每个损失调用 scaler.scale。
如果你的网络有多个优化器，你可以分别对其中任何一个调用 scaler.unscale_，
并且你必须分别对每个优化器调用 scaler.step。
然而，scaler.update 应该只在所有使用本次迭代的优化器都已更新后调用一次：
scaler = torch.cuda.amp.GradScaler()
for epoch in epochs:
for input, target in data:
optimizer0.zero_grad()
optimizer1.zero_grad()
with autocast(device_type='cuda', dtype=torch.float16):
output0 = model0(input)
output1 = model1(input)
loss0 = loss_fn(2 * output0 + 3 * output1, target)
loss1 = loss_fn(3 * output0 - 5 * output1, target)
# （这里的retain_graph与amp无关，它存在是因为在这个
# 示例中，两个backward()调用共享了图的某些部分。）
scaler.scale(loss0).backward(retain_graph=True)
scaler.scale(loss1).backward()
# 你可以选择哪些优化器接收显式的unscaling，如果你
# 想要检查或修改它们拥有的参数的梯度。
scaler.unscale_(optimizer0)
scaler.step(optimizer0)
scaler.step(optimizer1)
scaler.update()
每个优化器都会检查其梯度是否为无穷大/NaN，并独立决定是否跳过该步骤。这可能导致一个优化器跳过该步骤，而另一个优化器不跳过。由于跳过步骤的情况很少发生（每几百次迭代一次），这不应妨碍收敛。如果在向多优化器模型添加梯度缩放后观察到收敛不良，请报告错误。
使用多个GPU¶
这里描述的问题仅影响 autocast。GradScaler 的使用方式保持不变。
单个进程中的DataParallel¶
即使 torch.nn.DataParallel 生成线程以在每个设备上运行前向传递。
自动混合精度状态会在每个线程中传播，并且以下内容将正常工作：
model = MyModel()
dp_model = nn.DataParallel(model)
# 在主线程中设置自动转换
with autocast(device_type='cuda', dtype=torch.float16):
# dp_model的内部线程将自动转换。
output = dp_model(input)
# loss_fn也会自动转换
loss = loss_fn(output)
DistributedDataParallel，每个进程一个GPU¶
torch.nn.parallel.DistributedDataParallel的文档建议每个进程使用一个GPU以获得最佳性能。在这种情况下，DistributedDataParallel不会在内部生成线程，因此autocast和GradScaler的使用不受影响。
DistributedDataParallel，每个进程多个GPU¶
这里 torch.nn.parallel.DistributedDataParallel 可能会生成一个辅助线程在每个设备上运行前向传播，类似于 torch.nn.DataParallel。解决方法是一样的：
将 autocast 作为模型 forward 方法的一部分应用，以确保在辅助线程中启用它。
自动转换与自定义自动求导函数¶
如果你的网络使用了自定义自动求导函数
（torch.autograd.Function的子类），如果任何函数需要与自动混合精度兼容，则需要进行更改。
在所有情况下，如果你正在导入函数并且无法更改其定义，一个安全的回退方案是禁用自动转换并在使用时强制执行float32（或dtype），以避免在任何使用点发生错误：
with autocast(device_type='cuda', dtype=torch.float16):
...
with autocast(device_type='cuda', dtype=torch.float16, enabled=False):
output = imported_function(input1.float(), input2.float())
如果你是函数的作者（或可以修改其定义），更好的解决方案是使用
torch.cuda.amp.custom_fwd() 和 torch.cuda.amp.custom_bwd() 装饰器，如下面的相关案例所示。
具有多个输入或自动转换操作的函数¶
将 custom_fwd 和 custom_bwd（不带参数）分别应用于 forward 和
backward。这些确保 forward 在当前的自动转换状态下执行，并且 backward
在与 forward 相同的自动转换状态下执行（这可以防止类型不匹配错误）：
class MyMM(torch.autograd.Function):
@staticmethod
@custom_fwd
def forward(ctx, a, b):
ctx.save_for_backward(a, b)
return a.mm(b)
@staticmethod
@custom_bwd
def backward(ctx, grad):
a, b = ctx.saved_tensors
return grad.mm(b.t()), a.t().mm(grad)
现在可以在任何地方调用 MyMM，而无需禁用自动转换或手动转换输入：
mymm = MyMM.apply
with autocast(device_type='cuda', dtype=torch.float16):
output = mymm(input1, input2)
需要特定dtype的函数¶
考虑一个需要 torch.float32 输入的自定义函数。
应用 custom_fwd(cast_inputs=torch.float32) 到 forward
和 custom_bwd（不带参数）到 backward。
如果 forward 在启用了自动转换的区域内运行，装饰器会将浮点 CUDA Tensor 输入转换为 float32，并在 forward 和 backward 期间局部禁用自动转换：
class MyFloat32Func(torch.autograd.Function):
@staticmethod
@custom_fwd(cast_inputs=torch.float32)
def forward(ctx, input):
ctx.save_for_backward(input)
...
return fwd_output
@staticmethod
@custom_bwd
def backward(ctx, grad):
...
现在可以在任何地方调用 MyFloat32Func，无需手动禁用自动转换或转换输入：
func = MyFloat32Func.apply
with autocast(device_type='cuda', dtype=torch.float16):
# func 将以 float32 运行，无论周围的 autocast 状态如何
output = func(input)
