快捷導(dǎo)航

Pytorch中retain_graph的坑及解決

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

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

Pytorch中retain_graph的坑

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

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

?? ??? ?############################
? ? ? ? # (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

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

但是另外一個(gè)問(wèn)題在于，如果你都這么用的話，顯存會(huì)爆炸，因?yàn)樗Ａ袅颂荻?，所以都沒(méi)有及時(shí)釋放掉,浪費(fèi)資源。

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

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

Pytorch中的機(jī)制是每次調(diào)用loss.backward()時(shí)都會(huì)free掉計(jì)算圖中所有緩存的buffers，當(dāng)模型中可能有多次backward()時(shí)，因?yàn)榍耙淮握{(diào)用backward()時(shí)已經(jīng)釋放掉了buffer，所以下一次調(diào)用時(shí)會(huì)因?yàn)閎uffers不存在而報(bào)錯(cuò)