文章目錄

• PyTorch自動混合精度訓練(AMP)手冊
• • Autocasting
  • Gradient Scaling
  • Notes
  • Autocast Op Look-up-table
  • Reference

PyTorch自動混合精度訓練(AMP)手冊

自動混合精度 —— Automatic Mixed Precision, AMP

混合精度訓練是指在訓練過程中，一些操作使用float32資料類型的單精度，一些操作(linear/conv)使用float16資料類型的半精度。而自動混合精度訓練則是指，自動給每個操作比對其合适的資料類型(精度)。

一般來說，結合使用

torch.cuda.amp.autocast

和

torch.cuda.amp.GradScaler

，即可在torch中實作自動混合精度訓練。torch.cuda.amp.autocast會為選中的區域開啟autocasting功能，autocasting會自動為CUDA操作選擇(單/半)精度來提升性能并保持精度(實際上可能會帶來些許精度的提升)。torch.cuda.amp.GradScaler幫助執行梯度縮放步驟，梯度縮放會通過最小化梯度的underflow，來提升包含半精度(float16)梯度的網絡的收斂。

Autocasting

torch實作autocasting主要是通過API——

torch.cuda.amp.autocast

，autocast執行個體對象是作為上下文管理器(context manger)或裝飾器(decorator)來允許使用者代碼的某些區域在混合精度下運作，在這些區域中，CUDA操作會以一種由autocast選擇的特定于操作的資料類型運作，以此來提升性能并保持準确率。

autocast應該隻封裝網絡的前向傳播(forward pass(es))，以及損失計算(loss computation(s))。反向傳播不推薦在autocast區域内執行，反向傳播的操作會自動以對應的前向傳播的操作的資料類型運作。

代碼示例如下：

......
CUDA = torch.cuda.is_available()
device = torch.device("cuda" if CUDA else "cpu")
model = model.to(device)
optimizer = optim.SGD(model.parameters(), ...)
......

# Creates a GradScaler once at the begin
scaler = amp.GradScaler(enabled=CUDA)

for inputs, targets in data:
    optimizer.zero_grad()
    ......
    
    # Autocast
    # Enables autocasting for the forward pass (model + loss)
    with amp.autocast(enabled=CUDA):
        # Forward
        preds = model(inputs)
        loss = loss_fn(preds, targets)
        
    # Exits the context manager before backward()
    # Scales loss. Calls backward() on scaled loss to create scaled gradients.
    scaler.scale(loss).backward()
    
    # scaler.step() first unscales the gradients of the optimizer's assigned params. If these gradients do not contains infs or NaNs, optimizer.step() is then called. Otherwise, optimizer.step() is skipped.
    scaler.step(optimizer)
    
    # Updates the scale for next iteration.
    scaler.update()
           

Gradient Scaling

如果前向傳播對于一個特定的操作的輸入是float16資料類型的，那麼該操作對應的反向傳播也會産生float16的梯度。小幅值的梯度值可能在半精度下是不可表示的。這些值可能會重新整理為零(稱為underflow)，是以對應參數的更新也會丢失。(這裡可能是指類似梯度消失的問題，半精度下，如果梯度幅值小，權重在更新時幾乎沒有改變。)

為了避免underflow，采用**梯度縮放(gradient scaling)**的方式，将網絡的損失乘以一個縮放因子，并對縮放後的損失調用反向傳播，然後，通過網絡反向流動的梯度将按相同的因子進行縮放。也就是說，縮放後梯度值會有一個較大的幅值，是以它們不會被重新整理為零。

每個參數的梯度在優化器(optimizer)更新參數之前應該是未縮放的，是以縮放因子不會影響學習率。

代碼示例在上方，下方為api詳細說明。

scaler = torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)

# Params
# init_scale (float, optional, default=2.**16): 初始的縮放因子。
# growth_factor (float, optional, default=2.0): 如果在growth_interval個連續疊代中沒有出現inf/NaN的梯度，調用scaler.update()操作會執行growth_factor*scale操作，growth_factor主要用于增大縮放因子，一般為大于1.0的數。
# backoff_factor (float, optional, default=0.5): 如果在一個連續疊代中出現了inf/NaN的梯度，調用scaler.update()操作會執行growth_factor*scale操作，backoff_factor主要用于衰減縮放因子，一般為0.0~1.0的數。
# growth_interval (int, optional, default=2000): growth_interval用于控制在多少個連續疊代次數未出現inf/NaN的梯度時去增大縮放因子(即執行update():growth_factor*scale)。
# enabled (bool, optional, default=True): If False, disables gradient scaling. 'step' op simply invokes the underlying 'optimizer.step()', and other methods become no-ops. 

# Methods
# scale()方法會将輸入的一個張量或一個清單的張量乘以一個縮放因子。
scaler.scale(outputs) -> scaled outputs

# step()方法執行兩個操作：
# 1. 内部調用unscale_(optimizer)方法，除非之前已經對optimizer顯式調用過unscale_()方法了。作為unscale_()方法的部分功能，會檢查梯度是否為inf/NaN。
# 2. 如果沒有發現有inf/NaN的梯度，使用解縮放後的梯度調用optimizer.step(*args, **kwargs)。否則，會跳過optimizer.step()，以避免破壞參數。
scaler.step(optimizer, *args, **kwargs) -> the return value of optimizer.step(*args, **kwargs)

# unscale_()方法會讓優化器optimizer的梯度張量除以(devide，也可了解為unscale)縮放因子。
# unscale_()方法是可選的，一般用于需要在反向傳播跟step()更新參數之間修改(modify)或檢查(inspect)梯度的情況。如果沒有顯式調用，也會在調用scaler.step()時内部自動調用。
scaler.unscale_(optimizer)

# update()方法用于更新縮放因子。傳入new_scale參數來直接設定縮放因子scale。
# 如果任一優化器step操作被跳過了，則将縮放因子scale乘以backoff_factor來減小它；如果growth_intervel個連續疊代都沒有跳過優化器step操作，則将縮放因子scale乘以growth_factor來增大它。
# update()方法應該隻在疊代的最後被調用，即在一個疊代中給所有使用的優化器調用scaler.step(optimizer)操作後使用。
scaler.update(new_scale=None)

           

Notes

1. 隻有CUDA操作才适用于autocasting；運作在雙精度float64或非浮點資料類型下的操作不适用于autocasting，這兩種情況不會被執行autocasting，仍會運作在原資料類型下；隻有out-of-place操作(即非inplace操作)和Tensor方法才适用于autocasting；顯式指定dtype參數的操作op(…, dtype=xx)不适用于autocasting，此類操作會生成指定的dtype類型的輸出。
2. 反向傳播不推薦在autocast區域内執行。
3. 在autocast區域内生成的浮點張量(Floating-point Tensors)可能是半精度，float16資料類型的。在傳回非autocast區域時，如果将上述浮點張量與不同資料類型的浮點張量一起使用，可能導緻類型不比對錯誤(type mismatch errors)。在這種情況下，應将手動(manually)将autocast區域産生的半精度浮點張量強制轉換成單精度(或其他需要的資料類型)，如執行tensor.float()操作。
4. autocast區域内可以繼續嵌套autocast區域或非autocast區域，可通過autocast()的enabled參數控制，簡單示例如下。嵌套操作在需要強制某子區域要運作在特定資料類型下時非常有效，但仍需要注意避免張量類型不比對的問題。

   with autocast(enabled=True):
       >>>> op1 with autocast
       
       with autocast(enabled=False):
           >>>> op2 needed no autocast
          
       >>>> other ops
              

5. autocast狀态是線程本地(thread-local)的，意思是想要autocast在一個新的線程下運作，必須在該線程下調用上下文管理器或裝飾器。如果按基本的操作執行，即隻在執行主程序的main檔案中設定autocasting區域時，在多GPU運作(一個GPU一個程序)的DP或DDP模式下，autocasting将不會起任何作用。簡單了解：多GPU運作的DP/DDP模式會在每個device上生成線程來執行前向傳播操作，執行autocasting的是main檔案生成的主線程，而執行前向傳播操作的是各個device各自生成的線程，因為autocasting是線程本地(局域)作用的，是以即使在autocasting區域内，執行在其他線程的前向傳播操作不會被autocast。

   model = MyModel()
   dp_model = nn.DataParallel(model)
   
   # Sets autocast in the main thread
   with autocast():
       # dp_model's internal threads won't autocast.  The main thread's autocast state has no effect.
       output = dp_model(input)
       # loss_fn still autocasts, but it's too late...
       loss = loss_fn(output)
              

   解決辦法很簡單：在原設定的基礎上，将autocast融合進model的forward方法，要麼将autocast()作為forward()方法的裝飾器，要麼在forward()方法内部設定autocasting區域。(按如下示例操作，上方的代碼就會起作用了)

   MyModel(nn.Module):
       ...
       @autocast()
       def forward(self, input):
          ...
   
   # Alternatively
   MyModel(nn.Module):
       ...
       def forward(self, input):
           with autocast():
               ...
              

Autocast Op Look-up-table

Reference

• AUTOMATIC MIXED PRECISION PACKAGE - TORCH.CUDA.AMP
• AUTOMATIC MIXED PRECISION EXAMPLES
• PyTorch Recipes – AUTOMATIC MIXED PRECISION