从PyTorch中的梯度计算出发谈如何避免训练GAN中出现inplace error

从PyTorch中的梯度计算出发谈如何避免训练GAN中出现inplace error

one of the variables needed for gradient computation has been modified by an inplace operation这个错误在训练对抗网络时很容易出现，这往往是由于不熟悉PyTorch的计算图和梯度传播机制造成的。

叶子结点与非叶子结点

import torch
a = torch.tensor([1., 2, 3]).requires_grad_(True)
b = a * 2
loss = b.sum()
loss.backward()

PyTorch中自己创建的张量被称为叶子结点，叶子结点默认是不带梯度的，如果需要进行梯度计算，需要将其requires_grad属性设置为True， 而其它由叶子结点参与的运算所产生的张量都是非叶子结点。对于上面的例子而言，a就是叶子结点，b和loss都是非叶子结点。
我们只能获取叶子结点的梯度，而非叶子结点的梯度一般会在梯度回传后释放。

In [1]: import torch

In [2]: a = torch.tensor([1., 2, 3], requires_grad=True)

In [3]: b = torch.tensor([2. ,3, 4])

In [4]: c = a * b

In [5]: loss = c.sum()

In [6]: loss.backward()

In [7]: a.grad
Out[7]: tensor([2., 3., 4.])

In [8]: b.grad

In [9]: loss.grad
E:Anacondainstallscriptsipython:1: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. 
Its .grad attribute won't be populated during autograd.backward(). If you indeed want the gradient for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. 
If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more information.

观察上面的例子，变量a和变量b都是叶子结点，但是只有a的requires_grad属性被设置为True，所以a.grad可以得到本次运算的梯度，而b没有梯度。对于loss这个非叶子结点，我们想要访问它的梯度是被报警告的。

计算图

深度模型框架在实现梯度回传算法时需要先构建一个计算图，PyTorch中的计算图是动态的，梯度回传以后计算图会被释放。

所谓计算图就是在前向传播时构建了结点之间的运算关系，这样在梯度回传时就可以根据该计算图准确的计算出各个结点的梯度，方便后续进行参数更新。拿上图举例，当我们在loss进行梯度回传时，变量c，a相对于loss的梯度沿着前向传播的反方向依次被计算。
动态计算图的指的是，当叶子结点的梯度得到以后，前向传播构建的结点之间运算关系，以及非叶子结点本次梯度回传时计算得出的梯度都会被释放。这样及时的释放资源可以使得计算效率变高，这也是为什么PyTorch不允许我们访问非叶子结点的梯度，因为这些资源都被释放了。
基于上面的原因，loss也无法二次反向传播，因为第一次方向传播结束以后，各节点之间的运算关系已经被释放了，梯度无法回传。

inplace *** 作

inplace *** 作就是直接对变量的内容进行修改，而不是采用中间变量的方式去接收。

    a = torch.tensor([1, 2., 3], requires_grad=True)
    b = torch.tensor([2, 3, 4.])
    c = a * b
    a += 1
    loss = c.sum()
    loss.backward()
   Traceback (most recent call last):
File "E:/DAProject/DASS/config.py", line 30, in 
    a += 1
RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.

a += 1这种就是很典型的inplace *** 作。而我们一定要避免对叶子结点进行inplace *** 作，因为这会使得在反向传播时造成无法预期的后果。当我们的程序很简单时，程序会明确的指出错误位置，可一旦程序变得复杂，inplace *** 作就不好debug了。
关于inplace *** 作有两个注意点，第一点是要避免对叶子结点进行inplace *** 作，非叶子结点无所谓，所以我们在CNN网络中经常可以看到ReLU(inplace=True)的语句，这里是对非叶子结点进行的，而inplace *** 作又可以节省内存，这样使用一般是没有错的。

GAN中无法避免的inplace *** 作

在我们训练GAN的时候，我们似乎明明没有使用inplace *** 作，但是依然会报相关的错误，这是因为什么呢？
因为**optimizer.step()**就是个inplace *** 作且无法避免。optimizer.step会利用反向传播计算得到的梯度对神经网络的参数进行更新，而这个参数的更新也必然是在参数的内容上进行更改的。
下面通过一个简单的例子，对GAN训练的常见错误进行模拟，并说明如何避免类似的问题出现。

import torch
from torch import nn, optim
G = nn.Linear(2, 2)
D = nn.Linear(2, 2)
optimG = optim.Adam(G.parameters())
optimD = optim.Adam(D.parameters())

t = torch.tensor([1., 2], requires_grad=True)
g = G(t)
d = D(g)

lossD = d.sum()
optimD.zero_grad()
lossD.backward()
optimD.step()

lossG = (g + d).sum()
optimG.zero_grad()
lossG.backward()
optimG.step()

根据以上示意图我们来逐行分析代码。
我们先训练的是discriminator，它的loss是只依赖于自己的输出d。

lossD.backward()
optimD.step()

这两行代码是重中之重，lossD.backward()执行以后，D和G的参数的梯度都会被计算，与此同时，前向传播所构建的结点之间的关系都被释放了。而optimizerD.step()会根据刚刚计算的梯度对D中的参数进行更新，这是一个inplace *** 作。

再看下面训练generator的代码

lossG = (g + d).sum()
optimG.zero_grad()
lossG.backward()
optimG.step()

lossG的计算既依赖于G也依赖于D的输出。此时我们执行lossG.backward()一定会报错，因为从G到D的计算图都被释放了。

Traceback (most recent call last):
  File "E:/DAProject/DASS/config.py", line 19, in 
    lossG.backward()
  File "E:Anacondainstalllibsite-packagestorch_tensor.py", line 255, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "E:Anacondainstalllibsite-packagestorchautograd__init__.py", line 149, in backward
    allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag
RuntimeError: Trying to backward through the graph a second time (or directly access saved variables after they have already been freed).
Saved intermediate values of the graph are freed when you call .backward() or autograd.grad().
Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved variables after calling backward.

现在我们的需求是当lossD进行反向传播时保留前向传播的关系，这就需要retain_graph=True这个属性了

lossD.backward(retain_graph=True)

加了这个参数以后，计算图就不会被释放，lossG就能正常的反向传播。可是，程序此时依然有错。前面提到optimD.step()执行以后，D的参数被修改了且是inplace *** 作，而lossG的计算依赖于D的输出d，那么lossG在反向传播时肯定会发生inplace error。

Traceback (most recent call last):
  File "E:/DAProject/DASS/config.py", line 19, in 
    lossG.backward()
  File "E:Anacondainstalllibsite-packagestorch_tensor.py", line 255, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "E:Anacondainstalllibsite-packagestorchautograd__init__.py", line 149, in backward
    allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [2, 2]], 
which is output 0 of TBackward, is at version 2; expected version 1 instead.

现在是训练G的参数，D的参数应该是不变的，所以我们不希望D参数的梯度在反向传播时被计算，这样既节省了计算资源，又可以避免inplace error。

lossG = (g + d.detach()).sum()

detach()返回一个新的张量，这个张量和d共享数据，但是不参与梯度计算，用来起到截断梯度流的作用。
完整正确的代码如下所示：

import torch
from torch import nn, optim
G = nn.Linear(2, 2)
D = nn.Linear(2, 2)
optimG = optim.Adam(G.parameters())
optimD = optim.Adam(D.parameters())

t = torch.tensor([1., 2], requires_grad=True)
g = G(t)
d = D(g)

lossD = d.sum()
optimD.zero_grad()
lossD.backward(retain_graph=True)
optimD.step()

lossG = (g + d.detach()).sum()
optimG.zero_grad()
lossG.backward()
optimG.step()