Python PyTorch GradScaler.unscale_用法及代码示例

本文简要介绍python语言中 torch.cuda.amp.GradScaler.unscale_ 的用法。

用法:

unscale_(optimizer)

参数：

optimizer(torch.optim.Optimizer) -拥有未缩放梯度的优化器。

将 (“unscales”) 优化器的梯度张量除以比例因子。

unscale_() 是可选的，适用于需要修改或检查向后传递和 step() 之间的梯度的情况。如果未显式调用 unscale_()，则在 step() 期间将自动取消缩放渐变。

简单的示例，使用 unscale_() 启用未缩放渐变的裁剪：

...
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
scaler.step(optimizer)
scaler.update()

注意

unscale_() 不会产生 CPU-GPU 同步。

警告

每个优化器每次 step() 调用只能调用 unscale_() 一次，并且只能在该优化器分配的参数的所有梯度都已累积之后调用。在每个 step() 之间为给定优化器调用 unscale_() 两次会触发运行时错误。

警告

unscale_() 可能会取消缩放稀疏梯度，从而替换 .grad 属性。