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 屬性。