使用PyTorch的自動混合精度

教學

PyTorch

引言

更大的深度學習模型需要更多的計算能力和記憶資源。透過新技術的發展，深度神經網絡的訓練速度已經得到提升。而非使用FP32（完整精度浮點數格式），您可選擇使用FP16（半精度浮點數格式），研究者們發現二者結合使用是更好的選擇。

混精度允許在半精度訓練的同時，仍能保留大部分單精度網絡的準確性。”混精度技術”這個詞是指這種方法同時利用單精度與半精度表示。

在這個PyTorch下的自動混精度（AMP）訓練概述中，我們展示技術如何工作，並逐步通過使用AMP的過程，並讨论AMP技術更進階的應用，為用戶提供代碼支架以便後續與自身代碼整合。

前提知識

PyTorch基本知識：熟悉PyTorch，包括它如張量、模塊和訓練循環等核心概念。

深度學習基礎理解：如神經網絡、反向傳播和優化等概念。

混合精度訓練的知識: 了解混合精度訓練的優點和缺點，包括減少記憶體使用和加快計算速度。

访问兼容硬件: 支持混合精度的GPU，如具有张量核心的NVIDIA GPU（例如，Volta、Turing、Ampere架构）。

Python 和 CUDA 设置: 安装有PyTorch並為GPU加速配置好的Python環境。

混合精度概览

與大多數深度學習框架一樣，PyTorch通常在32位浮點數(FP32)上進行訓練。然而，FP32並不总是必要的。對於一些操作，可以使用16位浮點數，而FP32則會消耗更多時間和記憶體。

因此，NVIDIA工程師開發了一種技術，允許在進行大部分網絡運算時使用FP16，而對少部分操作則使用FP32的混合精度訓練。

在可行的情況下將模型轉換為使用浮點16(float16)數據类型。
保持浮點32(float32)的主權重以每隔一次迭代累計權重更新。
使用損失缩放來維持微小的梯度值。

對於混合精度訓練，PyTorch已經內建了許多功能。
模組的參數在呼叫.half()方法時會被轉為FP16，而在呼叫.half()時張量的數據會被轉為FP16。這些模組或張量上的任何運算都会使用快速的FP16算術執行。PyTorch很好地支援NVIDIA数学庫（cuBLAS和cuDNN）。FP16管線的數據會使用張量核心來執行GEMM和卷積。要在cuBLAS中使用張量核心，GEMM的尺寸([M, K] x [K, N] -> [M, N])必須是8的倍數。

介紹Apex

Apex的混合精度工具旨在增加訓練速度，同時保持單精度訓練的準確性和穩定性。Apex可以進行FP16或FP32的運算，自動處理主參數轉化，並自動缩放損失。

Apex 是為了讓研究人員更容易將混合精度訓練纳入其模型而創建的。Amp（ Automatic Mixed-Precision 的缩寫）是 Apex 的一個功能，是一個輕量级的 PyTorch 擴展。用戶只需要在他們的網絡上增加幾行代碼，就可以通過 Amp 享受到混合精度訓練的益處。Apex 是在 CVPR 2018 上推出的，值得注意的是，自推出以來，PyTorch 社區已經對 Apex 展示了強烈的支持。

通過對運行模型進行微小的更改，Amp 使得在創建或運行您的腳本時，您不需要擔心混合類型。Amp 的假設可能不太適合使用 PyTorch 的 uncommon 方法的模型，但是有挂钩可以根據需要調整這些假設。

Amp 提供了混合精度訓練的所有優點，而無需管理損失縮放或類型轉換。Apex 的 GitHub 網站涵蓋了安裝程序的指示，其官方 API 文档可以在這裡找到。

Amp 的運作方式

AMP 於邏輯層面上使用白名單/黑名單 paradigm。PyTorch中的張量操作包括例如torch.nn.functional.conv2d的神经網絡功能、例如torch.log的簡單數學功能以及例如torch.Tensor. add__的張量方法。這個宇宙中有三種主要類型的函數：

白名單：可能從 FP16 數學的效能提升中受益的函數。典型的應用包括矩陣乘法和卷積。
黑名單：對於精度可能不足的16位元，對於這些函數輸入應該是 FP32。
其他所有（剩餘的函數）：可以運行在 FP16 之中，但是由於效能提升不顯著，所以將其從 FP32 轉為 FP16 執行並不值得。

AMP 的任務在理論上很简单，至少。AMP 在呼叫之前確定一個 PyTorch 函數是否在白名單、黑名单或既不在白名单也不在黑名单。如果是在白名單，所有參數都應該轉為 FP16；如果是在黑名單，所有參數都應該轉為 FP32。如果既不在白名单也不在黑名单，只需確保所有參數的型態相同。這個政策在實際應用上比表面上看起來要複雜。

使用Amp與PyTorch模型結合

要在目前的PyTorch腳本中包含Amp，請按照以下步驟進行：

使用Apex庫導入Amp。
初始化Amp，使其可以更改模型、優化器以及PyTorch內部函數所需的設定。
注意反向傳播 (.backward()) 發生的地方，以便Amp可以同時縮放損失並清除每 iteration 狀態。

步驟1

第一步只需要一行代碼：

from apex import amp

步驟2

用於訓練的神经網絡模型和優化器必須已經指定，才能完成這一步，它只一行長。

model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

其他設定可讓您細調Amp的張量與操作類型的調整。函數 amp.initialize() 接受許多參數，我們只會指定其中的三個：

models (torch.nn.Module 或torch.nn.Modules的列表) – 需要修改/轉型的模型。
優化器 (選擇性, torch.optim.Optimizer 或 torch.optim.Optimizers 的列表) – 需要修改/转型的優化器。對於訓練是必須的，對於推理是選擇性的。
opt_level (str, 選擇性, 預設=“O1”) – 純或混合格式優化等級。接受的值有“O0”, “O1”, “O2”, 和 “O3”, 上面有詳細的解释。有四個優化等級:

O0 對於 FP32 訓練: 這是一個無操作。由於你的模型應該已經是 FP32 格式，因此無需擔心 O0，它可能有助於建立準確性的基準線。

O1 對於典型使用的混合格式 (建議使用): 修改所有 Tensor 和 Torch 方法，使其使用白名单-黑名單輸入變換方案。在 FP16 中，白名單操作如 GEMMs 和卷積等 Tensor Core 友善的運算會被执行。例如，Softmax 是一個需要 FP32 準確性的黑名單運算。除非另有說明，O1 也使用动态损失缩放。

O2 對於“几乎 FP16” 的混合格式: O2 將模型權重轉為 FP16，修补模型的前向方法以將輸入數據轉為 FP16，保留批次归一化在 FP32，保持 FP32 主导权重，更新優化器的 param_groups，使優化器的 step() 直接對 FP32 權重進行操作，並实施动态损失缩放（除非被覆蓋）。與 O1 不同，O2 不修补 Torch 函數或 Tensor 方法。

O3用於FP16訓練：O3在真正混合精度方面可能不如O1和O2穩定。因此，為您的模型設定一個基線速度可能是有益的，以便可以評估O1和O2的效率。
O3中額外的屬性覆蓋keep_batchnorm_fp32=True可能會幫助您確定如果您的模型使用批次正規化，“速度光”的效率，從而使cudnn批次正規化得以啟用。

O0和O3並不是真正的混合精度，但它们分别帮助设定准确性和速度基线。混合精度的实现被定义为O1和O2。
您可以嘗試兩者，看看哪一個對於您特定的模型來說提高性能和精度最多。

步驟3

確保您標記出您的代碼中反向傳播發生的地方。
以下幾行程式碼會出現：

loss = criterion(…)
loss.backward()
optimizer.step()

步驟4

使用Amp上下文管理器，您可以通過簡單地包裹反向傳播來啟用損失縮放：

loss = criterion(…)
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
optimizer.step()

就是這樣。現在您可以將混合精度訓練開啟並重新運行您的脚本。

捕捉函數調用

由於PyTorch非常靈活且動態，因此它缺乏固定的模型物件或圖路徑以固定的方式插入上面提到的轉換。通過對必要的函數進行“猴子补丁”修復，Amp能夠动态拦截並转换参数。

例如，您可以使用下面的代碼來確保torch.nn.functional.linear方法的參數总是被轉換為fp16:

orig_linear = torch.nn.functional.linear
def wrapped_linear(*args):
 casted_args = []
  for arg in args:
    if torch.is_tensor(arg) and torch.is_floating_point(arg):
      casted_args.append(torch.cast(arg, torch.float16))
    else:
      casted_args.append(arg)
  return orig_linear(*casted_args)
torch.nn.functional.linear = wrapped_linear

雖然Amp可能添加了一些细化以增加代碼的韧性，但調用Amp.init()實際上會導致猴子补丁插入所有相關的PyTorch函數中，從而正確地在運行時間進行参数轉換。

最小化转换

由于Amp保留所有参数转换的內部缓存並按需重用，因此每次迭代每個權重只會被转换一次FP32到FP16。在每次迭代中，反向傳播的上下文管理器告訴Amp何時清空缓存。

使用PyTorch進行自动 casting 和梯度缩放

“自動混合精度訓練”是指結合了torch.cuda.amp.autocast和torch.cuda.amp.GradScaler的技術。透過使用torch.cuda.amp.autocast，您可以為特定區域設定自動映射。自動映射會自動為GPU操作選擇精度，以優化效率同時保持準確度。

torch.cuda.amp.GradScaler實例使得執行梯度缩放步驟更加便捷。梯度縮放可以減少梯度下溢，幫助使用浮點16梯度的網絡達到更好的汇聚。

以下是一些代碼示例，展示如何在PyTorch中使用autocast()來獲得自動混合精度：

# 在預設精度下創建模型和優化器
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# 訓練開始時創建一個GradScaler。
scaler = GradScaler()

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

        # 使用autocast运行前向傳播。
        with autocast(device_type='cuda', dtype=torch.float16):
            output = model(input)
            loss = loss_fn(output, target)

        # 後向操作使用autocast為相對應前向操作選擇的精度。
        scaler.scale(loss).backward()

        # scaler.step()首先將優化器分配的參數的梯度進行解縮。
   
        scaler.step(optimizer)

        # 更新下一迭代的精度。
        scaler.update()

如果某個操作的前向傳遞具有float16輸入，那麼這個操作的反向傳遞會產生float16梯度，而float16可能無法表示小梯度大小。

如果這些值被冲淡水（”underflow”）時，相關參數的更新將會丢失。

梯度缩放是一種技術，使用縮放因子乘以網絡的損失，然後在縮放的損失上進行反向傳輸，以避免冲淡水。同時，也必須通過相同的因子縮減通過網絡的反向梯度。因此，梯度值具有較大的大小，防止它們被冲淡水。

在更新參數之前，每個參數的梯度（.grad屬性）應該被解縮，以確保縮放因子不會干擾學習率。因為它們是模塊化的，所以autocast和GradScaler都可以獨立使用。

與未缩放的梯度一同工作

梯度修剪

我們可以透過使用 `Scaler.scale(Loss).backward()` 方法來放大所有梯度。在 `backward()` 和 `scaler.step(optimizer)` 之間，`param.grad` 屬性必須在更改或檢查它們之前还被縮放。如果您想要限制梯度組的全球范數（請見torch.nn.utils.clip_grad_norm_()）或最大大大小（請見torch.nn.utils.clip_grad_value_()）至某個值以下，您可以使用一種稱為“梯度裁剪”的技術。

如果不取消縮放而裁剪，則梯度的范數/最大大小的縮放，使得您請求的門檻（原先是對於未縮放的梯度的門檻）無效。由優化器的給定参数包含的梯度通過 scaler.unscale (optimizer) 進行取消縮放。
您可以使用 scaler.unscale (optimizer1) 來取消縮放先前傳給其他優化器（如 optimizer1）的其他参数的梯度。我們可以通過添加兩行程式碼來说明這個概念：

# 在原地取消縮放優化器指定的params的梯度
        scaler.unscale_(optimizer)
# 由於優化器指定的params的梯度已被取消縮放，因此照常裁剪： 
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

與縮放渐變一起工作

漸變積累

漸變積累是基於一個荒謬的基本概念。並不是更新模型參數，而是等待並將 Gradients 積累通過連續的批次來計算損失和漸變。

在一定数量的批次之後，根據累計的漸變更新參數。以下是如何使用 漸變積累使用 pytorch 的代码片段：

scaler = GradScaler()

for epoch in epochs:
    for i, (input, target) in enumerate(data):
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
            # 正常化損失 
            loss = loss / iters_to_accumulate

        # 積累縮放的漸變。
        scaler.scale(loss).backward()
          # 權重更新
        if (i + 1) % iters_to_accumulate == 0:
            # 如果需要，可以在這裡進行 unscale_ 
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

漸變積累將渐變跨足 batch_per_iter * iters_to_accumulate 的适量批次大小。
缩放應該為有效批次進行校准；這意味著檢查是否存在 inf/NaN 成績，如果發現任何 inf/NaN，則跳過步驟，並將縮放更新到有效批次的精細度。
當為特定有效批次添加漸變時，也必須保持渐變在縮放和一致的縮放因子中。

如果学生在累加完成之前未被尺度化（或者比例因子发生变化），下一次反向传播将会把尺度化的学生与未被尺度化的学生（或者按不同因子尺度化的学生）相加，此后无法恢复累积的未被尺度化的学生的步伐必须适用。

你可以通过在每一步之前不久使用unscale对学生进行尺度化，在所有即将到来的步骤的尺度化学生被累积之后。unscale。
为了确保完整的有效批次，只需在每次迭代的末尾调用update，你之前在step
enumerate(data)函数允许我们在遍历数据时跟踪批次索引。
将运行损失除以iters_to_accumulate(loss / iters_to_accumulate)。这通过规范化损失来减少我们处理的每个迷你批次的影响。如果你在每个批次内平均损失，除法已经是正确的，不需要进一步规范化。这一步取决于你是如何计算损失的，可能是不必要的。
當我們使用 `scaler.scale(loss).backward()` 時，PyTorch 会将缩放的梯度累积起来，直到我们调用 `optimizer.zero_grad()` 为止。

梯度惩罚

在實現梯度惩罚時，使用 `torch.autograd.grad()` 來建立梯度，這些梯度會合併形成惩罚值，然後加到損失中。以下示例展示了沒有缩放或自動 casting 的 L2 惩罚。

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)

        # 创建梯度
        grad_prams = torch.autograd.grad(outputs=loss,
                                          inputs=model.parameters(),
                                          create_graph=True)

        # 计算惩罚项并加到损失中
        grad_norm = 0
        for grad in grad_prams:
            grad_norm += grad.pow(2).sum()
        grad_norm = grad_norm.sqrt()
        loss = loss + grad_norm

        loss.backward()

        # 您可以在這裡裁剪梯度

        optimizer.step()

提供給 `torch.autograd.grad()` 的张量应该缩放到实施梯度惩罚。在合并以获得惩罚值之前，需要反缩放梯度。由于惩罚项的計算是前向傳输的一部分，因此应该在自動 casting 上下文中進行。
對於相同的 L2 惩罚，以下是它的样子：

scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        #為autograd.grad的反向傳播執行損失縮放，結果#scaled_grad_prams
        scaled_grad_prams = torch.autograd.grad(outputs=scaler.scale(loss),
                                                 inputs=model.parameters(),
                                                 create_graph=True)

        #在計算懲罰之前創建grad_prams（grad_prams必須是#未縮放的）。 
        #因為沒有優化器擁有scaled_grad_prams，所以使用 conventioal除法而不是scaler.unscale_:
        inv_scaled = 1./scaler.get_scale()
        grad_prams = [p * inv_scaled for p in scaled_grad_prams]

        #計算並將懲罰項添加到損失中。 
        with autocast():
            grad_norm = 0
            for grad in grad_prams:
                grad_norm += grad.pow(2).sum()
            grad_norm = grad_norm.sqrt()
            loss = loss + grad_norm

        #對反向調用應用縮放。
        #正確縮放葉子梯度。
        scaler.scale(loss).backward()

        #您可以在這裡進行unscale_ 

        #步驟()和更新()像往常一樣進行。
        scaler.step(optimizer)
        scaler.update()

與多個模型、損失和優化器一起工作

Scaler.scale 必须在您的網絡中的每個損失被調用，如果您有很多的話。
如果您在您的網絡中有很多優化器，您可以在它們之中的任何一个上调用 scaler.unscale，並且您必須在它們每一个上调用 scaler.step。然而，scaler.update 應該只使用一次，在所有在此迭代中使用的優化器進行 steppping 之後：

scaler = torch.cuda.amp.GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer1.zero_grad()
        optimizer2.zero_grad()
        with autocast():
            output1 = model1(input)
            output2 = model2(input)
            loss1 = loss_fn(2 * output1 + 3 * output2, target)
            loss2 = loss_fn(3 * output1 - 5 * output2, target)

       #雖然 retain graph 与 amp 無關，但它出現在這個 #示例中，因為兩個 backward() 調用共享某些圖的区域。 
        scaler.scale(loss1).backward(retain_graph=True)
        scaler.scale(loss2).backward()

        #如果您希望查看或調整它們拥有的梯度，您可以指定哪些優化器進行显式的不縮放。.
        scaler.unscale_(optimizer1)

        scaler.step(optimizer1)
        scaler.step(optimizer2)

        scaler.update()

每個優化器檢查其梯度是否存在 infs/NaNs，並個別判断是否跳過步驟。一些優化器可能會跳過步驟，而其他則可能不會。步驟跳過發生在每几百次迭代一次；因此，它應該不會影響 convergence。對於多優化器模型，如果您在添加梯度縮放後看到 convergence 欠佳，您可以報告問題。

與多張GPU合作

深度學習模型的一個重要問題是，它們變得太大，無法在單張GPU上進行訓練。在單張GPU上訓練一個模型的時間可能會太長，需要多張GPU同時訓練以盡可能地加快準備模型的速度。一位知名研究者能夠將ImageNet的訓練時間從兩週缩短到18分鐘，或者在兩週內訓練最廣泛和最先进的Transformer-XL，而不是需要四年的時間。

DataParallel 和 DistributedDataParallel

quality], PyTorch 提供了最優的易用性與控制性的組合。nn.DataParallel 和 nn.parallel.DistributedDataParallel 是 PyTorch 用於在多個 GPU 上分發訓練的兩個特性。您可以使用這些易於使用的封裝器和更改，在多個 GPU 上訓練網絡。

单个进程中的 DataParallel

在單個機器上，DataParallel 幫助將訓練分發到很多 GPU 上。
讓我們近距離觀察 DataParallel 在實際中的应用方式。
當使用 DataParallel 訓練神經網絡時，以下階段会发生：

来源：

迷你批次在 GPU:0 上分割。
分割並分發迷你批次到所有可用的 GPU 上。
將模型複製到 GPU 上。
在所有 GPU 上進行前向傳遞。
在 GPU:0 上計算與網絡輸出相關的損失，並將損失返回給各種 GPU。各 GPU 上應該計算梯度。
GPU上梯度的總和：0，並將優化器應用於更新模型。

需要注意的是，這裡 Discussed concerns only apply to autocast. GradScaler 的行為保持不變。是否 torch.nn.DataParallel 为每个设备创建线程进行前向传播，并不重要。autocast 状态在每一个中进行沟通，以下将会工作：

model = Model_m()
p_model = nn.DataParallel(model)

 # 在主线程中设置 autocast
with autocast():
     # p_model 中将有 autocasting 
    output = p_model(input)
     # loss_fn 也将 autocast
    loss = loss_fn(output)

每个进程一个 GPU 的 DistributedDataParallel

torch.nn.parallel.DistributedDataParallel的文件說明，為了獲得最佳性能，建議每個程式的使用一個GPU。在這種情況下，DistributedDataParallel internal並不會啟動線程；因此autocast和GradScaler的使用并不会受到影响。

DistributedDataParallel, 每個程式使用多個GPU

在torch.nn.parallel.DistributedDataParallel中，可能會為每個設備創建一個旁路線程來運行前向傳播，就像torch.nn.DataParallel一樣。解決方案是相同的：將autocast作為模型前向方法的一部分應用，以確保在旁路線程中启用它。

結論

在本文中，我們介紹了：

介绍了Apex。
了解Amps如何工作。
了解如何進行梯度缩放、梯度修剪、梯度累积和梯度罰款。
了解如何與多個模型、損失函數和優化器一起工作。
了解在處理多個GPU時，如何在一個进程中進行DataParallel。

參考資料

https://developer.nvidia.com/blog/apex-pytorch-easy-mixed-precision-training/
diy7>https://nvidia.github.io/apex/amp.html
https://discuss.pytorch.org/t/accumulating-gradients/30020
https://towardsdatascience.com/how-to-scale-training-on-multiple-gpus-dae1041f49d2

Source:
https://www.digitalocean.com/community/tutorials/automatic-mixed-precision-using-pytorch