PyTorch訓(xùn)練LSTM時loss.backward()報錯的解決方案

訓(xùn)練用PyTorch編寫的LSTM或RNN時，在loss.backward()上報錯：

RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time.

千萬別改成loss.backward(retain_graph=True)，會導(dǎo)致顯卡內(nèi)存隨著訓(xùn)練一直增加直到OOM：

RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 10.73 GiB total capacity; 9.79 GiB already allocated; 13.62 MiB free; 162.76 MiB cached)

正確做法：

LSRM / RNN模塊初始化時定義好hidden，每次forward都要加上self.hidden = self.init_hidden():
Class LSTMClassifier(nn.Module):
    def __init__(self, embedding_dim, hidden_dim)：
    # 此次省略其它代碼
    self.rnn_cell = nn.LSTM(embedding_dim, hidden_dim)
    self.hidden = self.init_hidden()
    # 此次省略其它代碼
    
    def init_hidden(self):
        # 開始時刻, 沒有隱狀態(tài)
        # 關(guān)于維度設(shè)置的詳情,請參考 Pytorch 文檔
        # 各個維度的含義是 (Seguence, minibatch_size, hidden_dim)
        return (torch.zeros(1, 1, self.hidden_dim),
                torch.zeros(1, 1, self.hidden_dim))
    def forward(self, x):
        # 此次省略其它代碼
        self.hidden = self.init_hidden()  # 就是加上這句!!!!
        out, self.hidden = self.rnn_cell(x, self.hidden)     
        # 此次省略其它代碼
        return out    

或者其它模塊每次調(diào)用這個模塊時，其它模塊的forward()都對這個LSTM模塊init_hidden()一下。

如定義一個模型LSTM_Model()：

Class LSTM_Model(nn.Module):
    def __init__(self, embedding_dim, hidden_dim)：
        # 此次省略其它代碼
        self.rnn = LSTMClassifier(embedding_dim, hidden_dim)
        # 此次省略其它代碼
        
    def forward(self, x):
        # 此次省略其它代碼
        self.rnn.hidden = self.rnn.init_hidden()  # 就是加上這句!!!!
        out = self.rnn(x)     
        # 此次省略其它代碼
        return out

這是因為：

根據(jù) 官方tutorial，在 loss 反向傳播的時候，pytorch 試圖把 hidden state 也反向傳播，但是在新的一輪 batch 的時候 hidden state 已經(jīng)被內(nèi)存釋放了，所以需要每個 batch 重新 init （clean out hidden state）， 或者 detach，從而切斷反向傳播。

補(bǔ)充：pytorch:在執(zhí)行l(wèi)oss.backward()時out of memory報錯

在自己編寫SurfNet網(wǎng)絡(luò)的過程中，出現(xiàn)了這個問題，查閱資料后，將得到的解決方法匯總?cè)缦?/p>

可試用的方法：

1、reduce batch size, all the way down to 1

2、remove everything to CPU leaving only the network on the GPU

3、remove validation code, and only executing the training code

4、reduce the size of the network (I reduced it significantly: details below)

5、I tried scaling the magnitude of the loss that is backpropagating as well to a much smaller value

在訓(xùn)練時，在每一個step后面加上：

torch.cuda.empty_cache()

在每一個驗證時的step之后加上代碼：

with torch.no_grad()

不要在循環(huán)訓(xùn)練中累積歷史記錄

total_loss = 0
for i in range(10000):
    optimizer.zero_grad()
    output = model(input)
    loss = criterion(output)
    loss.backward()
    optimizer.step()
    total_loss += loss

total_loss在循環(huán)中進(jìn)行了累計，因為loss是一個具有autograd歷史的可微變量。你可以通過編寫total_loss += float(loss)來解決這個問題。

本人遇到這個問題的原因是，自己構(gòu)建的模型輸入到全連接層中的特征圖拉伸為1維向量時太大導(dǎo)致的，加入pool層或者其他方法將最后的卷積層輸出的特征圖尺寸減小即可。

以上為個人經(jīng)驗，希望能給大家一個參考，也希望大家多多支持腳本之家。