快捷導(dǎo)航

Pytorch中retain_graph的坑及解決

更新時間：2023年02月21日 09:08:40 作者：Longlongaaago

這篇文章主要介紹了Pytorch中retain_graph的坑及解決方案，具有很好的參考價值，希望對大家有所幫助。如有錯誤或未考慮完全的地方，望不吝賜教

Pytorch中retain_graph的坑

在查看SRGAN源碼時有如下?lián)p失函數(shù)，其中設(shè)置了retain_graph=True，其作用就是

在更新D網(wǎng)絡(luò)時的loss反向傳播過程中使用了retain_graph=True，目的為是為保留該過程中計算的梯度，后續(xù)G網(wǎng)絡(luò)更新時使用；

?? ??? ?############################
? ? ? ? # (1) Update D network: maximize D(x)-1-D(G(z))
? ? ? ? ###########################
? ? ? ? real_img = Variable(target)
? ? ? ? if torch.cuda.is_available():
? ? ? ? ? ? real_img = real_img.cuda()
? ? ? ? z = Variable(data)
? ? ? ? if torch.cuda.is_available():
? ? ? ? ? ? z = z.cuda()
? ? ? ? fake_img = netG(z)
?
? ? ? ? netD.zero_grad()
? ? ? ? real_out = netD(real_img).mean()
? ? ? ? fake_out = netD(fake_img).mean()
? ? ? ? d_loss = 1 - real_out + fake_out
? ? ? ? d_loss.backward(retain_graph=True) #####
? ? ? ? optimizerD.step()
?
? ? ? ? ############################
? ? ? ? # (2) Update G network: minimize 1-D(G(z)) + Perception Loss + Image Loss + TV Loss
? ? ? ? ###########################
? ? ? ? netG.zero_grad()
? ? ? ? g_loss = generator_criterion(fake_out, fake_img, real_img)
? ? ? ? g_loss.backward()
? ? ? ? optimizerG.step()
? ? ? ? fake_img = netG(z)
? ? ? ? fake_out = netD(fake_img).mean()
?
? ? ? ? g_loss = generator_criterion(fake_out, fake_img, real_img)
? ? ? ? running_results['g_loss'] += g_loss.data[0] * batch_size
? ? ? ? d_loss = 1 - real_out + fake_out
? ? ? ? running_results['d_loss'] += d_loss.data[0] * batch_size
? ? ? ? running_results['d_score'] += real_out.data[0] * batch_size
? ? ? ? running_results['g_score'] += fake_out.data[0] * batch_size

也就是說，只要我們有一個loss，我們就可以先loss.backward(retain_graph=True) 讓它先計算梯度，若下面還有其他損失，但是可能你想擴展代碼，可能有些loss是不用的，所以先加了 if 等判別語句進行了干預(yù)，使用loss.backward(retain_graph=True)就可以單獨的計算梯度，屢試不爽。

但是另外一個問題在于，如果你都這么用的話，顯存會爆炸，因為他保留了梯度，所以都沒有及時釋放掉,浪費資源。

而正確的做法應(yīng)該是，在你最后一個loss 后面，一定要加上loss.backward()這樣的形式，也就是讓最后一個loss 釋放掉之前所有暫時保存下來得梯度?。?/p>

Pytorch中有多次backward時需要retain_graph參數(shù)

Pytorch中的機制是每次調(diào)用loss.backward()時都會free掉計算圖中所有緩存的buffers，當模型中可能有多次backward()時，因為前一次調(diào)用backward()時已經(jīng)釋放掉了buffer，所以下一次調(diào)用時會因為buffers不存在而報錯