前言

本文讲解基于pytorch0.4.0版本，如不清楚版本信息请看这里。backward()在pytorch中是一个经常出现的函数，我们一般会在更新loss的时候使用它，比如loss.backward()。通过对loss进行backward来实现从输出到输入的自动求梯度运算。但是这里的backward()如果追根溯源一下，或者说Go to definition一下，我们会发现，其实这个backward是来源于torch.autograd.backward。

上面是官方的截图信息。但是这个函数我们可能不常见，那么这个函数在哪儿呢，就在Tensor这个类中(之前是在Variable类中，现在Variable和tensor合并)。而Tensor这个类中有一个函数：

backward()函数，这个函数返回的就是torch.autograd.backward()。也就是说，我们在训练中输入我们数据，然后经过一系列神经网络运算，最后计算loss，然后loss.backward()。这里的backward()归根绝地就是，上面说的这个函数。

正文

本文要说明的两个backward，一个nn.Module中的backward()和torch.autograd.Function中的backward()，其实有一个是假的backward()。

Fake Backward

很容易发现，我们在自己定义一个全新的网络层的时候会继承nn.Module，但是我们只需要实现__init__和forward()即可，不需要实现也没必要实现backward()函数，即使你实现了，你继承了nn.Module并且编写了一个backward()函数：

class ContentLoss(nn.Module):

    def __init__(self, target, weight):
        super(ContentLoss, self).__init__()
        # we 'detach' the target content from the tree used
        self.target = target.detach() * weight
        # to dynamically compute the gradient: this is a stated value,
        # not a variable. Otherwise the forward method of the criterion
        # will throw an error.
        self.weight = weight
        self.criterion = nn.MSELoss()

    def forward(self, input):
        self.loss = self.criterion(input * self.weight, self.target)
        self.output = input
        return self.output

    def backward(self, retain_graph=True):
        print('ContentLoss Backward works')
        self.loss.backward(retain_graph=retain_graph)
        return self.loss

    ...
# 执行backward语句，具体代码请看下方的连接。
for sl in style_losses:
    style_score += sl.backward()
for cl in content_losses:
    content_score += cl.backward()

上面这段代码是利用pytorch实现风格迁移的自定义内容损失层，如果不懂看这里：传送门。如果正常操作，在实际运行中上面的backward函数并不会执行也不回打印执行信息。上面定义的backward称为fake backward函数，也就是假的backward函数，不会在pytorch的自动求梯度图中执行。但是为什么这么写，在pytorch官方0.3.0的教程中，可以在loss更新的时候，不使用loss.backward(),而是直接使用类中的.backward()方法然后返回loss即可。

但是在官方的0.4.0的风格迁移示例代码中，上面的代码发生了变化：

class ContentLoss(nn.Module):

    def __init__(self, target, ):
        super(ContentLoss, self).__init__()
        # we 'detach' the target content from the tree used
        # to dynamically compute the gradient: this is a stated value,
        # not a variable. Otherwise the forward method of the criterion
        # will throw an error.
        self.target = target.detach()

    def forward(self, input):
        self.loss = F.mse_loss(input, self.target)
        return input

...
# 执行代码，具体看官网的最新0.4.0风格迁移教程
for sl in style_losses:
    style_score += sl.loss
for cl in content_losses:
    content_score += cl.loss

loss = style_score + content_score
loss.backward()

我们发现没有backward函数了，而且使用的loss function发生了变化，从nn.MSELoss() ==> F.mse_loss()。

上面的这段代码没有定义backward函数，也没有执行retain_grad操作。为什么两个版本的不一样，其实第一个版本(0.3.0)完全没必要写backward函数，也没必要再单独执行backward()函数，因为最终目的都是一样的，都是要实现对loss的backward，在forward中进行操作的时候，其实我们已经对torch.autograd.Function的subclass进行了操作。也就是说在我们对tensor进行每一步操作运算的时候都会生成一个Function类的子类，里面定了好了forward和backward操作，最后连成计算图，所以没有必要多此一举。

说了这么多，既然不建议在nn.Module中定义backward。那我们能不能自己定义backward函数。

Real Backward

可以的。

通过继承torch.autograd.Function来定义。这一方面官方有教程，这里就不赘述。（下方是官方示例程序）

class MyReLU(torch.autograd.Function):
  """
  We can implement our own custom autograd Functions by subclassing
  torch.autograd.Function and implementing the forward and backward passes
  which operate on Tensors.
  """
  @staticmethod
  def forward(ctx, x):
    """
    In the forward pass we receive a context object and a Tensor containing the
    input; we must return a Tensor containing the output, and we can use the
    context object to cache objects for use in the backward pass.
    """
    ctx.save_for_backward(x)
    return x.clamp(min=0)

  def backward(ctx, grad_output):
    """
    In the backward pass we receive the context object and a Tensor containing
    the gradient of the loss with respect to the output produced during the
    forward pass. We can retrieve cached data from the context object, and must
    compute and return the gradient of the loss with respect to the input to the
    forward function.
    """
    x, = ctx.saved_tensors
    grad_x = grad_output.clone()
    grad_x[x < 0] = 0
    return grad_x

这里讲一下我们在什么情况下需要自己定义：

我们平常使用的nn.Module其实说白了就是一层包装(Contain)，比如nn.Conv2继承了nn.Module，但是里面的核心函数是torch.nn.function.conv2d，为什么要包装下，原因很简单，为了方便，因为我们使用的卷积层是有参数的，这些参数是可以学习的(learnable parameters)。在这个包装类中我们通过torch.nn.parameter的Parameter类把参数进行包装然后传递给torch.nn.function中的函数进行计算，这样也就简化了我们的操作。

那么什么时候需要使用torch.autograd.Function去定义自己的层，在有些操作通过组合pytorch中已有的层实现不了的时候，比如你要实现一个新的梯度下降算法，那么就可以尝试着写这些东西。但是要注意，因为这个涉及到了底层，你需要forward和backward一起写，然后自己写对中间变量的操作，比如gradinput以及gradoutput。

比如这样写：

class my_function(Function):
    def forward(self, input, parameters):
        self.saved_for_backward = [input, parameters]
        # output = [对输入和参数进行的操作，这里省略]
        return output

    def backward(self, grad_output):
        input, parameters = self.saved_for_backward
        # grad_input = [求 forward(input)关于 parameters 的导数] * grad_output
        return grad_input

# 然后通过定义一个Module来包装一下

class my_module(nn.Module):
    def __init__(self, ...):
        super(my_module, self).__init__()
        self.parameters = # 初始化一些参数

    def forward(self, input):
        output = my_function(input, self.parameters) # 在这里执行你之前定义的function!
        return output

这样你就可以通过自定义层然后包装，然后来使用了。

后记

对于这个包装，其实包不包装对于执行效率的影响几乎可以不计，对于没有学习参数的层，比如Relu（nn.ReLU vs F.relu），其实包不包装对于使用起来没什么区别。

参考链接

https://discuss.pytorch.org/t/defining-backward-function-in-nn-module/5047

https://discuss.pytorch.org/t/whats-the-difference-between-torch-nn-functional-and-torch-nn/681

https://discuss.pytorch.org/t/difference-of-methods-between-torch-nn-and-functional/1076

https://discuss.pytorch.org/t/whats-the-difference-between-torch-nn-functional-and-torch-nn/681/4