PyTorch中梯度為什么默認自動累加,在反向傳播前要手動將梯度清零?


主要可以通過兩個角度來說明,但其實表述的意思也是異曲同工

低顯存跑大batchsize的角度

這種模式可以讓梯度玩出更多花樣,比如說梯度累加(gradient accumulation)

傳統的訓練函數,一個batch是這么訓練的:

for i,(images,target) in enumerate(train_loader):
    # 1. input output
    images = images.cuda(non_blocking=True)
    target = torch.from_numpy(np.array(target)).float().cuda(non_blocking=True)
    outputs = model(images)
    loss = criterion(outputs,target)

    # 2. backward
    optimizer.zero_grad()   # reset gradient
    loss.backward()
    optimizer.step()            
  1. 獲取loss:輸入圖像和標簽,通過infer計算得到預測值,計算損失函數;
  2. optimizer.zero_grad() 清空過往梯度;
  3. loss.backward() 反向傳播,計算當前梯度;
  4. optimizer.step() 根據梯度更新網絡參數

簡單的說就是進來一個batch的數據,計算一次梯度,更新一次網絡

 

使用梯度累加是這么寫的:

for i,(images,target) in enumerate(train_loader):
    # 1. input output
    images = images.cuda(non_blocking=True)
    target = torch.from_numpy(np.array(target)).float().cuda(non_blocking=True)
    outputs = model(images)
    loss = criterion(outputs,target)

    # 2.1 loss regularization
    loss = loss/accumulation_steps   
    # 2.2 back propagation
    loss.backward()
    # 3. update parameters of net
    if((i+1)%accumulation_steps)==0:
        # optimizer the net
        optimizer.step()        # update parameters of net
        optimizer.zero_grad()   # reset gradient
  1. 獲取loss:輸入圖像和標簽,通過infer計算得到預測值,計算損失函數;
  2. loss.backward() 反向傳播,計算當前梯度;
  3. 多次循環步驟1-2,不清空梯度,使梯度累加在已有梯度上;
  4. 梯度累加了一定次數后,先optimizer.step() 根據累計的梯度更新網絡參數,然后optimizer.zero_grad() 清空過往梯度,為下一波梯度累加做准備;

總結來說:梯度累加就是,每次獲取1個batch的數據,計算1次梯度,梯度不清空,不斷累加,累加一定次數后,根據累加的梯度更新網絡參數,然后清空梯度,進行下一次循環。

一定條件下,batchsize越大訓練效果越好,梯度累加則實現了batchsize的變相擴大,如果accumulation_steps為8,則batchsize '變相' 擴大了8倍,是我們這種乞丐實驗室解決顯存受限的一個不錯的trick,使用時需要注意,學習率也要適當放大。

 

 

方便多任務學習的角度

而這樣的好處可以從內存消耗的角度來看

1. Edition1

在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的來源組成更加復雜,內存消耗會更大

2. Edition2

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

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的設計理念出發的。

 

以上兩個解釋轉自知乎靠前的熱評回答,具體可參見https://www.zhihu.com/question/303070254


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM