PyTorch 的混合精度訓(xùn)練主要由兩個(gè)方法實(shí)現(xiàn)：amp.autocast 和 amp.GradScalar。在這兩個(gè)工具的幫助下，可以實(shí)現(xiàn)以 torch.float16 的混合精度訓(xùn)練。當(dāng)然，這兩個(gè)方法都是模塊化并且通常都會(huì)一起調(diào)用，但并不一定總是需要一起使用。

參考：

Automatic Mixed Precision package - torch.amp
Automatic Mixed Precision examples
Automatic Mixed Precision recipe

一 amp.autocast

在 PyTorch 中，autocast 是一個(gè)用于自動(dòng)混合精度訓(xùn)練的上下文管理器。它的主要作用是根據(jù)操作的類型自動(dòng)選擇使用 float16（半精度）或 float32（單精度）進(jìn)行計(jì)算，從而提高訓(xùn)練速度并減少內(nèi)存使用。以下是 autocast 的具體功能和工作原理：

1. 自動(dòng)選擇精度
   操作類型：在深度學(xué)習(xí)中，不同的操作對數(shù)值精度的要求不同。例如，某些操作（如矩陣乘法）在 float16 下可以安全地執(zhí)行，而其他操作（如某些歸一化或損失計(jì)算）可能需要 float32 來保持?jǐn)?shù)值穩(wěn)定性。
   自動(dòng)化：使用 autocast 時(shí)，PyTorch 會(huì)自動(dòng)判斷每個(gè)操作的最佳精度，并在適當(dāng)?shù)那闆r下使用 float16，在其他情況下使用 float32。這使得開發(fā)者不需要手動(dòng)管理每個(gè)操作的精度，從而簡化了代碼。
2. 上下文管理器
   使用方式：autocast 通常用作上下文管理器，使用 with torch.cuda.amp.autocast(): 語句包裹需要進(jìn)行混合精度計(jì)算的代碼塊。在這個(gè)代碼塊內(nèi)，所有的張量操作都會(huì)根據(jù) autocast 的規(guī)則自動(dòng)選擇精度。
3. 性能提升
   加速訓(xùn)練：通過使用 float16 進(jìn)行計(jì)算，autocast 可以顯著提高訓(xùn)練速度，尤其是在支持半精度計(jì)算的 GPU 上（如 NVIDIA 的 Volta 和 Ampere 架構(gòu)）。減少內(nèi)存使用：使用 float16 還可以減少顯存的占用，使得更大的模型或更大的批量大小成為可能。
4. 示例
   一般來說，建議在模型前向傳遞和 loss計(jì)算中使用 autocast，而不建議在反向傳播和參數(shù)更新過程中使用，例如：

# 來源：https://pytorch.org/docs/stable/amp.html#torch.autocast
# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)for input, target in data:optimizer.zero_grad()# Enables autocasting for the forward pass (model + loss)with torch.autocast(device_type="cuda"):output = model(input)　# 前向傳遞loss = loss_fn(output, target) # loss計(jì)算# Exits the context manager before backward()loss.backward() # 后向傳遞（計(jì)算梯度）optimizer.step() # 參數(shù)更新

5. 總結(jié)
   autocast 是 PyTorch 中實(shí)現(xiàn)混合精度訓(xùn)練的關(guān)鍵工具，它通過自動(dòng)選擇操作的精度來提高訓(xùn)練速度和減少內(nèi)存使用，同時(shí)保持?jǐn)?shù)值穩(wěn)定性。使用 autocast 可以讓開發(fā)者更專注于模型的設(shè)計(jì)和訓(xùn)練，而不必?fù)?dān)心每個(gè)操作的精度管理。

二 amp.GradScalar

GradScaler 是 PyTorch 中用于混合精度訓(xùn)練的一個(gè)重要組件，主要用于處理梯度的縮放，以確保在使用半精度（float16）進(jìn)行訓(xùn)練時(shí)的數(shù)值穩(wěn)定性。以下是對 GradScaler 的詳細(xì)介紹：

1. 背景
   在混合精度訓(xùn)練中，使用 float16 進(jìn)行計(jì)算可以提高速度和減少內(nèi)存使用，但由于 float16 的數(shù)值范圍和精度較低，可能會(huì)導(dǎo)致梯度下溢（即梯度變得過小而被視為零）或上溢（即梯度變得過大而無法表示）。GradScaler 的作用就是解決這個(gè)問題。
2. 主要功能
   動(dòng)態(tài)縮放：GradScaler 會(huì)在反向傳播時(shí)動(dòng)態(tài)調(diào)整損失值的縮放因子，以避免梯度的數(shù)值不穩(wěn)定。它會(huì)根據(jù)當(dāng)前的訓(xùn)練狀態(tài)自動(dòng)選擇合適的縮放因子。
   防止下溢和上溢：通過將損失值乘以一個(gè)縮放因子，GradScaler 可以確保計(jì)算出的梯度不會(huì)因?yàn)閿?shù)值過小而消失，也不會(huì)因?yàn)閿?shù)值過大而溢出。
3. 使用流程
   使用 GradScaler 的典型流程如下：

初始化：創(chuàng)建 GradScaler 的實(shí)例;

scaler = torch.cuda.amp.GradScaler()

前向傳播：在前向傳播時(shí)，使用 autocast 來自動(dòng)選擇精度，不建議使用GradScaler；

   with torch.cuda.amp.autocast():output = model(inputs)loss = compute_loss(output, targets)

反向傳播:

• 使用 scaler.scale(loss).backward() 進(jìn)行反向傳播。這里，scale 方法會(huì)將損失值縮放，并計(jì)算梯度。
• 使用 scaler.step(optimizer) 更新模型參數(shù)。此步驟會(huì)將縮放后的梯度應(yīng)用到優(yōu)化器中。
• 使用 scaler.update() 更新縮放因子，以便在下一次迭代中使用。

4. 示例代碼

以下是一個(gè)使用 GradScaler 的簡單示例：

import torch
from torch.cuda.amp import GradScaler, autocastmodel = ...  # 初始化模型
optimizer = ...  # 初始化優(yōu)化器
scaler = GradScaler()  # 創(chuàng)建 GradScaler 實(shí)例for data, target in dataloader:optimizer.zero_grad()  # 清零梯度with autocast():  # 自動(dòng)選擇精度output = model(data) # 前向傳遞loss = compute_loss(output, target) # 計(jì)算lossscaler.scale(loss).backward()  # 縮放損失并反向傳播以計(jì)算梯度scaler.step(optimizer)  # 更新模型參數(shù)scaler.update()  # 更新縮放因子

5. 總結(jié)
   GradScaler 是 PyTorch 中實(shí)現(xiàn)混合精度訓(xùn)練的關(guān)鍵工具，它通過動(dòng)態(tài)縮放損失值來確保在使用 float16 進(jìn)行訓(xùn)練時(shí)的數(shù)值穩(wěn)定性。使用 GradScaler 可以有效地避免梯度下溢和上溢的問題，從而提高訓(xùn)練的可靠性和效率。通過結(jié)合 autocast 和 GradScaler，開發(fā)者可以在不犧牲模型性能的情況下，充分利用混合精度訓(xùn)練的優(yōu)勢。