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)樗Ａ袅颂荻?，所以都沒有及時(shí)釋放掉,浪費(fèi)資源。

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

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ò)

解決辦法

loss.backward(retain_graph=True)

錯(cuò)誤使用

• optimizer.zero_grad() 清空過(guò)往梯度；
• loss1.backward(retain_graph=True) 反向傳播，計(jì)算當(dāng)前梯度；
• loss2.backward(retain_graph=True) 反向傳播，計(jì)算當(dāng)前梯度；
• optimizer.step() 根據(jù)梯度更新網(wǎng)絡(luò)參數(shù)

因?yàn)槊看握{(diào)用bckward時(shí)都沒有將buffers釋放掉，所以會(huì)導(dǎo)致內(nèi)存溢出，迭代越來(lái)越慢（因?yàn)樘荻榷急４媪?，沒有free)

正確使用

• optimizer.zero_grad() 清空過(guò)往梯度；
• loss1.backward(retain_graph=True) 反向傳播，計(jì)算當(dāng)前梯度；
• loss2.backward() 反向傳播，計(jì)算當(dāng)前梯度；
• optimizer.step() 根據(jù)梯度更新網(wǎng)絡(luò)參數(shù)

最后一個(gè) backward() 不要加 retain_graph 參數(shù)，這樣每次更新完成后會(huì)釋放占用的內(nèi)存，也就不會(huì)出現(xiàn)越來(lái)越慢的情況了

總結(jié)

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