Pytorch 中retain_graph的坑

最新推荐文章于 2023-07-21 21:39:10 发布

DeepGoAI

最新推荐文章于 2023-07-21 21:39:10 发布

阅读量3.1k

点赞数

分类专栏： pytorch 文章标签： pytorch backward retain_graph

本文链接：https://blog.csdn.net/Willen_/article/details/89394766

版权

pytorch 专栏收录该内容

17 篇文章 0 订阅

订阅专栏

Pytorch 中retain_graph的坑

在查看SRGAN源码时有如下损失函数，其中设置了retain_graph=True,其作用就是

在更新D网络时的loss反向传播过程中使用了retain_graph=True，目的为是为保留该过程中计算的梯度，后续G网络更新时使用；


		############################
        # (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 等判别语句进行了干预，使用loss.backward(retain_graph=True)就可以单独的计算梯度，屡试不爽。

但是另外一个问题在于，如果你都这么用的话，显存会爆炸，因为他保留了梯度，所以都没有及时释放掉,浪费资源。

而正确的做法应该是，在你最后一个loss 后面，一定要加上loss.backward()这样的形式，也就是让最后一个loss 释放掉之前所有暂时保存下来得梯度！！