pytorch训练GAN时的detach()

　　我最近在学使用Pytorch写GAN代码，发现有些代码在训练部分细节有略微不一样，其中有的人用到了detach()函数截断梯度流，有的人没用detch()，取而代之的是在损失函数在反向传播过程当中将backward(retain_graph=True)，本文经过两个 gan 的代码，介绍它们的做用，并分析，不一样的更新策略对程序效率的影响。

　　这两个 GAN 的实现中，有两种不一样的训练策略：

• 先训练判别器（discriminator），再训练生成器（generator），这是原始论文Generative Adversarial Networks 中的算法
• 先训练generator，再训练discriminator

　　为了减小网络垃圾，GAN的原理网上一大堆，我这里就不重复赘述了，想要详细了解GAN原理的朋友，能够参考我专题文章：神经网络结构：生成式对抗网络（GAN）。

须要了解的知识：

　　detach()：截断node反向传播的梯度流，将某个node变成不须要梯度的Varibale，所以当反向传播通过这个node时，梯度就不会从这个node往前面传播。

更新策略

　　咱们直接下面进入本文正题，即，在 pytorch 中，detach 和 retain_graph 是干什么用的？本文将借助三段 GAN 的实现代码，来举例介绍它们的做用。

先训练判别器，再训练生成器

策略一

咱们分析循环中一个 step 的代码：

valid = torch.Tensor(imgs.size(0), 1).fill_(1.0).to(device)  # 真实标签，都是1
fake = torch.Tensor(imgs.size(0), 1).fill_(0.0).to(device)  # 假标签，都是0

# ########################
#        训练判别器       #
# ########################
real_imgs = imgs.to(device)     # 真实图片
z = torch.randn((imgs.shape[0], 100)).to(device)  # 噪声

gen_imgs = generator(z)  # 从噪声中生成假数据
pred_gen = discriminator(gen_imgs)  # 判别器对假数据的输出
pred_real = discriminator(real_imgs)  # 判别器对真数据的输出

optimizer_D.zero_grad()  # 把判别器中全部参数的梯度归零
real_loss = adversarial_loss(pred_real, valid)  # 判别器对真实样本的损失
fake_loss = adversarial_loss(pred_gen, fake)  # 判别器对假样本的损失
d_loss = (real_loss + fake_loss) / 2  # 两项损失相加取平均

# 下面这行代码十分重要，将在正文着重讲解
d_loss.backward(retain_graph=True)  # retain_graph=True 十分重要，不然计算图内存将会被释放
optimizer_D.step()  # 判别器参数更新

# ########################
#        训练生成器       #
# ########################
g_loss = adversarial_loss(pred_gen, valid)  # 生成器的损失函数
optimizer_G.zero_grad()  # 生成器参数梯度归零
g_loss.backward()  # 生成器的损失函数梯度反向传播
optimizer_G.step()  # 生成器参数更新

代码讲解

　　鉴别器的损失函数d_loss是由real_loss和fake_loss组成的，而fake_loss又是noise通过generator来的。这样一来咱们对d_loss进行反向传播，不只会计算discriminator 的梯度还会计算generator 的梯度（虽然这一步optimizer_D.step()只更新 discriminator 的参数），所以下面在更新generator参数时，要先将generator参数的梯度清零，避免受到discriminator loss 回传过来的梯度影响。

　　generator 的 损失在回传时，一样要通过 discriminator 网络才能传递回自身（系统从输入噪声到 Discriminator 输出，从头至尾只有一次前向传播，而有两次反向传播，故在第一次反向传播时，鉴别器要设置 backward(retain graph=True)，保持计算图不被释放。由于 pytorch 默认一个计算图只计算一次反向传播，反向传播后，这个计算图的内存就会被释放，因此用这个参数控制计算图不被释放。所以，在回传梯度时，一样也计算了一遍 discriminator 的参数梯度，只不过此次 discriminator 的参数不更新，只更新 generator 的参数，即 optimizer_G.step()。同时，咱们看到，下一个 step 首先将 discriminator 的梯度重置为 0，就是为了防止 generator loss 反向传播时顺带计算的梯度对其形成影响（还有上一步 discriminator loss 回传时累积的梯度）。

　　综上，咱们看到，为了完成一步参数更新，咱们进行了两次反向传播，第一次反向传播为了更新 discriminator 的参数，但多余计算了 generator 的梯度。第二次反向传播为了更新 generator 的参数，可是计算了 discriminator 的梯度，所以在写一个step，须要当即清零discriminator梯度。

　　若是你实在看不懂，就照着这个形式写代码就好了，反正形式都帮大家写好了。

策略二

　　这种策略我遇到的比较多，也是先训练鉴别器，再训练生成器

　　鉴别器训练阶段，noise 从 generator 输入，输出 fake data，而后 detach 一下，随着 true data 一块儿输入 discriminator，计算 discriminator 损失，并更新 discriminator 参数。生成器训练阶段，把没通过 detach 的 fake data 输入到discriminator 中，计算 generator loss，再反向传播梯度，更新 generator 的参数。这种策略，计算了两次 discriminator 梯度，一次 generator 梯度。感受这种比较符合先更新 discriminator 的习惯。缺点是，以前的 generator 生成的计算图得保留着，直到 discriminator 更新完，再释放。

valid = torch.Tensor(imgs.size(0), 1).fill_(1.0).to(device)  # 真实标签，都是1
fake = torch.Tensor(imgs.size(0), 1).fill_(0.0).to(device)  # 假标签，都是0

# ########################
#        训练判别器       #
# ########################
real_imgs = imgs.to(device)     # 真实图片
z = torch.randn((imgs.shape[0], 100)).to(device)  # 噪声

gen_imgs = generator(z)  # 从噪声中生成假数据
pred_gen = discriminator(gen_imgs.detach())  # 假数据detach()，判别器对假数据的输出
pred_real = discriminator(real_imgs)  # 判别器对真数据的输出

optimizer_D.zero_grad()  # 把判别器中全部参数的梯度归零
real_loss = adversarial_loss(pred_real, valid)  # 判别器对真实样本的损失
fake_loss = adversarial_loss(pred_gen, fake)  # 判别器对假样本的损失
d_loss = (real_loss + fake_loss) / 2  # 两项损失相加取平均

# 下面这行代码十分重要，将在正文着重讲解
d_loss.backward()  # retain_graph=True 十分重要，不然计算图内存将会被释放
optimizer_D.step()  # 判别器参数更新

# ########################
#        训练生成器       #
# ########################
g_loss = adversarial_loss(pred_gen, valid)  # 生成器的损失函数
optimizer_G.zero_grad()  # 生成器参数梯度归零
g_loss.backward()  # 生成器的损失函数梯度反向传播
optimizer_G.step()  # 生成器参数更新

先训练生成器，再训练判别器

 咱们分析循环中一个 step 的代码：

valid = Variable(Tensor(imgs.size(0), 1).fill_(1.0), requires_grad=False)  # 真实样本的标签，都是 1
fake = Variable(Tensor(imgs.size(0), 1).fill_(0.0), requires_grad=False)  # 生成样本的标签，都是 0
z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))  # 噪声
real_imgs = Variable(imgs.type(Tensor))     # 真实图片

# ########################
#        训练生成器       #
# ########################
optimizer_G.zero_grad()  # 生成器参数梯度归零
gen_imgs = generator(z)  # 根据噪声生成虚假样本
g_loss = adversarial_loss(discriminator(gen_imgs), valid)  # 用真实的标签+假样本，计算生成器损失
g_loss.backward()  # 生成器梯度反向传播，反向传播通过了判别器，故此时判别器参数也有梯度
optimizer_G.step()  # 生成器参数更新，判别器参数虽然有梯度，可是这一步不能更新判别器

# ########################
#        训练判别器       #
# ########################
optimizer_D.zero_grad()  # 把生成器损失函数梯度反向传播时，顺带计算的判别器参数梯度清空
real_loss = adversarial_loss(discriminator(real_imgs), valid)  # 真样本+真标签：判别器损失
fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)  # 假样本+假标签：判别器损失
d_loss = (real_loss + fake_loss) / 2  # 判别器总的损失函数
d_loss.backward()  # 判别器损失回传
optimizer_D.step()  # 判别器参数更新

　　为了更新生成器参数，用生成器的损失函数计算梯度，而后反向传播，传播图中通过了判别器，根据链式法则，不得不顺带计算一下判别器的参数梯度，虽然在这一步不会更新判别器参数。反向传播事后，noise 到 fake image 再到 discriminator 的输出这个前向传播的计算图就被释放掉了，后面也不会再用到。

　　接着更新判别器参数，此时注意到，咱们输入判别器的是两部分，一部分是真实数据，另外一部分是生成器的输出，也就是假数据。注意观察细节，在判别器前向传播过程，输入的假数据被 detach 了，detach 的意思是，这个数据和生成它的计算图“脱钩”了，即梯度传到它那个地方就停了，再也不继续往前传播（实际上也不会再往前传播了，由于 generator 的计算图在第一次反向传播事后就被释放了）。所以，判别器梯度反向传播，就到它本身身上为止。

　　所以，比起第一种策略，这种策略要少计算一次 generator 的全部参数的梯度，同时，也没必要刻意保存一次计算图，占用没必要要的内存。

　　但须要注意的是，在第一种策略中，noise 从 generator 输入，到 discriminator 输出，只经历了一次前向传播，discriminator 端的输出，被用了两次，一次是计算 discriminator 的损失函数，另外一次是计算 generator 的损失函数。

　　而在第这种策略中，noise 从 generator 输入，到discriminator 输出，计算 generator 损失，回传，这一步更新了 generator 的参数，并释放了计算图。下一步更新 discriminator 的参数时，generator 的输出通过 detach 后，又经过了一遍 discriminator，至关于，generator 的输出先后两次经过了 discriminator ，获得相同的输出。显然，这也是冗余的。

总结

综上，这两段代码各有利弊：

　　第一段代码，好处是 noise 只进行了一次前向传播，缺点是，更新 discriminator 参数时，多计算了一次 generator 的梯度，同时，第一次更新 discriminator 须要保留计算图，保证算 generator loss 时计算图不被销毁。

　　第三段代码，好处是经过先更新 generator ，使更新后的前向传播计算图能够放心被销毁，所以不用保留计算图占用内存。同时，在更新 discriminator 的时候，也不会像上面的那段代码，计算冗余的 generator 的梯度。缺点是，在 discriminator 上，对 generator 的输出算了两次前向传播，第二次又产生了新的计算图（但比第一次的小）。

一个多计算了一次 generator 梯度，一个多计算一次 discriminator 前向传播。所以，二者差异不大。若是 discriminator 比generator 复杂，那么应该采起第一种策略，若是 discriminator 比 generator 简单，那么应该采起第三种策略，一般状况下，discriminator 要比 generator 简单，故若是效果差很少尽可能采起第三种策略。

　　可是第三种先更新generator，再更新 discriminator 老是给人感受怪怪得，由于 generator 的更新须要 discriminator 提供准确的 loss 和 gradient，不然岂不是在瞎更新？

　　可是策略三，立刻用完立刻释放。综合来讲，仍是策略三最好，策略二其次，策略一最差（差在多计算一次 generator gradient 上，而一般多计算一次 generator gradient 的运算量比多计算一次 discriminator 前向传播的运算量大），所以，detach 仍是颇有必要的。

参考

Pytorch: detach 和 retain_graph

使用PyTorch进行GAN训练时对于梯度截断的思考.detach()