使用基本 PyTorch 的混合精度訓練

你將使用低精度浮點數型別來加速語言翻譯模型的訓練。例如，16 位元浮點數（float16）的大小只有 32 位元（float32）的一半。這能加速矩陣相乘與卷積等運算。請記得，這會涉及對梯度進行縮放，以及將運算轉為 16 位元浮點數。

以下物件已預先載入：dataset、model、dataloader、optimizer。

本練習屬於課程

使用 PyTorch 高效訓練 AI 模型

檢視課程

練習說明

在迴圈之前，使用 torch.amp.GradScaler 為梯度建立一個 scaler。
在迴圈中，使用 torch.autocast 作為情境管理器，將運算轉為 16 位元浮點數型別。
在迴圈中，對 loss 進行縮放並呼叫 .backward()，以建立縮放後的梯度。

動手互動練習

試著完成這個範例程式碼，體驗一下這個練習。

# Define a scaler for the gradients
scaler = torch.amp.____()
for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]
    # Casts operations to mixed precision
    with torch.____(device_type="cpu", dtype=torch.____):
        outputs = model(inputs, labels=targets)
        loss = outputs.loss
    # Compute scaled gradients
    scaler.____(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

編輯並執行程式碼