pytorch中的梯度累加(Gradient Accumulation)

PyTorch中，在反向傳播前為什么要手動將梯度清零？

原因在於，在PyTorch中，計算得到的梯度值會進行累加,而這樣的好處，可以從內存消耗的角度來看。
在PyTorch中，multi-task任務一個標准的train from scratch流程為:

for idx, data in enumerate(train_loader):
    xs, ys = data
    pred1 = model1(xs)
    pred2 = model2(xs)
    

    loss1 = loss_fn1(pred1, ys)
    loss2 = loss_fn2(pred2, ys)
    
    ******
    loss = loss1 + loss2
    optmizer.zero_grad()
    loss.backward()
    ++++++
    optmizer.step()

從PyTorch的設計原理上來說，在每次進行前向計算得到pred時，會產生一個用於梯度回傳的計算圖，這張圖儲存了進行back propagation需要的中間結果，當調用了.backward()后，會從內存中，將這張圖進行釋放。上述代碼執行到******時，內存中是包含了兩張計算圖的，而隨着求和得到loss，這兩張圖進行了合並，而且大小的變化可以忽略。執行到++++++時，得到對應的grad值並且釋放內存。這樣，訓練時必須存儲兩張計算圖，而如果loss的來源組成更加復雜，內存消耗會更大。

為了減小每次的內存消耗，借助梯度累加，又有,有如下變種:

for idx, data in enumerate(train_loader):
    xs, ys = data
    
    optmizer.zero_grad()
    # 計算d(l1)/d(x)
    pred1 = model1(xs) #生成graph1
    loss1 = loss_fn1(pred1, ys)
    loss1.backward()  #釋放graph1

    # 計算d(l2)/d(x)
    pred2 = model2(xs)#生成graph2
    loss2 = loss_fn2(pred2, ys)
    loss2.backward()  #釋放graph2

    # 使用d(l1)/d(x)+d(l2)/d(x)進行優化
    optmizer.step()

借助梯度累加，避免同時計算多個損失時,存儲多個計算圖。可以從代碼中看出，利用梯度累加，可以在最多保存一張計算圖的情況下，進行multi-task任務的訓練。另外一個理由就是，在內存大小不夠的情況下，疊加多個batch的grad，作為一個大batch進行迭代，因為二者得到的梯度是等價的綜上可知，這種梯度累加的思路是對內存的極大友好，是由FAIR的設計理念出發的。

梯度累加(Gradient Accumulation)

我們在訓練神經網絡的時候，超參數batch size的大小會對最終的模型效果產生很大的影響。一定條件下，batch size設置的越大，模型就會越穩定。batch size的值通常設置在 8-32 之間，但是當我們做一些計算量需求大的任務(例如語義分割、GAN等)或者輸入圖片尺寸太大的時候，我們的batch size往往只能設置為2或者4，否則就會出現 “CUDA OUT OF MEMORY” 的不可抗力報錯。貧窮是促進人類進步的階梯，如何在有限的計算資源的條件下，訓練時采用更大的batch size呢？這就是梯度累加(Gradient Accumulation)技術了。
我們以Pytorch為例，一個神經網絡的訓練過程通常如下：

for i, (inputs, labels) in enumerate(trainloader):
    optimizer.zero_grad()                   # 梯度清零
    outputs = net(inputs)                   # 正向傳播
    loss = criterion(outputs, labels)       # 計算損失
    loss.backward()                         # 反向傳播，計算梯度
    optimizer.step()                        # 更新參數
    if (i+1) % evaluation_steps == 0:
        evaluate_model()

從代碼中可以很清楚地看到神經網絡是如何做到訓練的：

1. 將前一個batch計算之后的網絡梯度清零
2. 正向傳播，將數據傳入網絡，得到預測結果
3. 根據預測結果與label，計算損失值
4. 利用損失進行反向傳播，計算參數梯度
5. 利用計算的參數梯度更新網絡參數

下面來看梯度累加是如何做的：

for i, (inputs, labels) in enumerate(trainloader):
    outputs = net(inputs)                   # 正向傳播
    loss = criterion(outputs, labels)       # 計算損失函數
    loss = loss / accumulation_steps        # 損失標准化
    loss.backward()                         # 反向傳播，計算梯度
    if (i+1) % accumulation_steps == 0:
        optimizer.step()                    # 更新參數
        optimizer.zero_grad()               # 梯度清零
        if (i+1) % evaluation_steps == 0:
            evaluate_model()

1. 正向傳播，將數據傳入網絡，得到預測結果
2. 根據預測結果與label，計算損失值
3. 利用損失進行反向傳播，計算參數梯度
4. 重復1-3，不清空梯度，而是將梯度累加
5. 梯度累加達到固定次數之后，更新參數，然后將梯度清零

總結來講，梯度累加就是每計算一個batch的梯度，不進行清零，而是做梯度的累加，當累加到一定的次數之后，再更新網絡參數，然后將梯度清零。通過這種參數延遲更新的手段，可以實現與采用大batch size相近的效果。在平時的實驗過程中，我一般會采用梯度累加技術，大多數情況下，采用梯度累加訓練的模型效果，要比采用小batch size訓練的模型效果要好很多。
一定條件下，batch size越大,訓練效果越好，梯度累加則實現了batch size的變相擴大，如果 accumulation_steps 為8，則batch size '變相' 擴大了8倍，是解決顯存受限的一個不錯的trick，使用時需要注意，學習率也要適當放大。
參考：
PyTorch中在反向傳播前為什么要手動將梯度清零？
PyTorch中的梯度累加
梯度累加(Gradient Accumulation)