自動混合精度示例

創建於：2020年2月13日 | 最後更新於：2024年9月13日

通常，“自動混合精度訓練”意味著將 torch.autocast 和 torch.amp.GradScaler 結合使用進行訓練。

torch.autocast 的例項允許在選定的區域內啟用自動型別轉換。自動型別轉換會自動選擇操作的精度，以提高效能同時保持準確性。

torch.amp.GradScaler 的例項有助於方便地執行梯度縮放的步驟。梯度縮放透過最小化梯度下溢來改善具有 float16（預設在 CUDA 和 XPU 上）梯度的網路的收斂性，具體解釋請參閱 此處。

torch.autocast 和 torch.amp.GradScaler 是模組化的。在下面的示例中，每個都按照其單獨的文件建議使用。

(這裡的示例僅供說明。有關可執行的詳細教程，請參閱 自動混合精度指南。)

典型的混合精度訓練

# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# Creates a GradScaler once at the beginning of training.
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # Runs the forward pass with autocasting.
        with autocast(device_type='cuda', dtype=torch.float16):
            output = model(input)
            loss = loss_fn(output, target)

        # Scales loss.  Calls backward() on scaled loss to create scaled gradients.
        # Backward passes under autocast are not recommended.
        # Backward ops run in the same dtype autocast chose for corresponding forward ops.
        scaler.scale(loss).backward()

        # scaler.step() first unscales the gradients of the optimizer's assigned params.
        # If these gradients do not contain infs or NaNs, optimizer.step() is then called,
        # otherwise, optimizer.step() is skipped.
        scaler.step(optimizer)

        # Updates the scale for next iteration.
        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()

        # Unscales the gradients of optimizer's assigned params in-place
        scaler.unscale_(optimizer)

        # Since the gradients of optimizer's assigned params are unscaled, clips as usual:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

        # optimizer's gradients are already unscaled, so scaler.step does not unscale them,
        # although it still skips optimizer.step() if the gradients contain infs or NaNs.
        scaler.step(optimizer)

        # Updates the scale for next iteration.
        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 呼叫之前呼叫 unscale_，此時將為即將進行的 step 呼叫累積完所有（已縮放的）梯度。另外，僅在您呼叫 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

        # Accumulates scaled gradients.
        scaler.scale(loss).backward()

        if (i + 1) % iters_to_accumulate == 0:
            # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)

            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)

        # Creates gradients
        grad_params = torch.autograd.grad(outputs=loss,
                                          inputs=model.parameters(),
                                          create_graph=True)

        # Computes the penalty term and adds it to the loss
        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()

        # clip gradients here, if desired

        optimizer.step()

要實現*帶有*梯度縮放的梯度懲罰，傳遞給 torch.autograd.grad() 的 outputs 張量（Tensor）應該被縮放。因此，生成的梯度也將被縮放，在組合起來建立懲罰值之前應該先進行縮放。

此外，懲罰項的計算是前向傳遞的一部分，因此應該在 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)

        # Scales the loss for autograd.grad's backward pass, producing scaled_grad_params
        scaled_grad_params = torch.autograd.grad(outputs=scaler.scale(loss),
                                                 inputs=model.parameters(),
                                                 create_graph=True)

        # Creates unscaled grad_params before computing the penalty. scaled_grad_params are
        # not owned by any optimizer, so ordinary division is used instead of scaler.unscale_:
        inv_scale = 1./scaler.get_scale()
        grad_params = [p * inv_scale for p in scaled_grad_params]

        # Computes the penalty term and adds it to the loss
        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

        # Applies scaling to the backward call as usual.
        # Accumulates leaf gradients that are correctly scaled.
        scaler.scale(loss).backward()

        # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)

        # step() and update() proceed as usual.
        scaler.step(optimizer)
        scaler.update()

處理多個模型、損失和最佳化器

如果您的網路有多個損失，您必須分別對每個損失呼叫 scaler.scale。如果您的網路有多個最佳化器，您可以分別對任何一個最佳化器呼叫 scaler.unscale_，並且您必須分別對每個最佳化器呼叫 scaler.step。

但是，scaler.update 應該只調用一次，在所有本次迭代使用的最佳化器都已執行 `step` 之後。

scaler = torch.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 here is unrelated to amp, it's present because in this
        # example, both backward() calls share some sections of graph.)
        scaler.scale(loss0).backward(retain_graph=True)
        scaler.scale(loss1).backward()

        # You can choose which optimizers receive explicit unscaling, if you
        # want to inspect or modify the gradients of the params they own.
        scaler.unscale_(optimizer0)

        scaler.step(optimizer0)
        scaler.step(optimizer1)

        scaler.update()

每個最佳化器都會檢查其梯度是否包含 inf/NaN，並獨立決定是否跳過步驟。這可能導致一個最佳化器跳過步驟，而另一個則不跳過。由於跳過步驟的情況很少見（每幾百次迭代一次），這不應該影響收斂。如果您在為多最佳化器模型新增梯度縮放後觀察到收斂不佳，請報告 bug。

處理多 GPU

此處描述的問題僅影響 autocast。 GradScaler 的使用保持不變。

單程序中的 DataParallel

即使 torch.nn.DataParallel 啟動執行緒在每個裝置上執行前向傳遞。自動型別轉換狀態會在每個執行緒中傳播，並且以下程式碼可以正常工作：

model = MyModel()
dp_model = nn.DataParallel(model)

# Sets autocast in the main thread
with autocast(device_type='cuda', dtype=torch.float16):
    # dp_model's internal threads will autocast.
    output = dp_model(input)
    # loss_fn also autocast
    loss = loss_fn(output)

DistributedDataParallel，每個程序一個 GPU

torch.nn.parallel.DistributedDataParallel 的文件建議每個程序使用一個 GPU 以獲得最佳效能。在這種情況下，DistributedDataParallel 不會在內部啟動執行緒，因此 autocast 和 GradScaler 的使用不受影響。

DistributedDataParallel，每個程序多個 GPU

在這種情況下，torch.nn.parallel.DistributedDataParallel 可能會啟動一個輔助執行緒在每個裝置上執行前向傳遞，類似於 torch.nn.DataParallel。 解決方法相同：將自動型別轉換作為模型 forward 方法的一部分應用，以確保它在輔助執行緒中啟用。

Autocast 和自定義 Autograd 函式

如果您的網路使用 自定義 autograd 函式（torch.autograd.Function 的子類），則需要進行更改以實現自動型別轉換相容性，如果任何函式

• 接受多個浮點張量輸入，

• 包裝任何可自動型別轉換的操作（請參閱 Autocast 操作參考），或

• 需要特定的 dtype（例如，如果它包裝了僅為 dtype 編譯的 CUDA 擴充套件）。

在所有情況下，如果您匯入了函式且無法修改其定義，一個安全的折衷方法是在出現錯誤的使用點停用自動型別轉換，並強制以 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.amp.custom_fwd() 和 torch.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(device_type='cuda', cast_inputs=torch.float32) 應用於 forward，並將 custom_bwd(device_type='cuda') 應用於 backward。如果 forward 在啟用了自動型別轉換的區域執行，這些裝飾器會將浮點張量輸入轉換為 float32，在由 device_type 引數指定的裝置上（在此示例中為 CUDA），並且在 forward 和 backward 執行期間本地停用自動型別轉換。

class MyFloat32Func(torch.autograd.Function):
    @staticmethod
    @custom_fwd(device_type='cuda', cast_inputs=torch.float32)
    def forward(ctx, input):
        ctx.save_for_backward(input)
        ...
        return fwd_output
    @staticmethod
    @custom_bwd(device_type='cuda')
    def backward(ctx, grad):
        ...

現在 MyFloat32Func 可以在任何地方呼叫，而無需手動停用自動型別轉換或轉換輸入。

func = MyFloat32Func.apply

with autocast(device_type='cuda', dtype=torch.float16):
    # func will run in float32, regardless of the surrounding autocast state
    output = func(input)