pytorch relu函数实现_Pytorch 自定义Function

转载：https://zhuanlan.zhihu.com/p/27783097

https://zhuanlan.zhihu.com/XavierLin

0. 本章内容

在本次，我们将学习如何自定义一个torch.autograd.Function，下面是本次的主要内容

1. 对Function的直观理解；

2. Function与Module的差异与应用场景；

3. 写一个简单的ReLU Function；

1.对Function的直观理解

在之前的介绍中，我们知道，Pytorch是利用Variable与Function来构建计算图的。回顾下Variable，Variable就像是计算图中的节点，保存计算结果(包括前向传播的激活值，反向传播的梯度)，而Function就像计算图中的边，实现Variable的计算，并输出新的Variable。Function简单说就是对Variable的运算，如加减乘除，relu，pool等。但它不仅仅是简单的运算。与普通Python或者numpy的运算不同，Function是针对计算图，需要计算反向传播的梯度。因此他不仅需要进行该运算(forward过程)，还需要保留前向传播的输入(为计算梯度)，并支持反向传播计算梯度。如果有做过公开课cs231的作业，记得里面的每个运算都定义了forward，backward，并通过保存cache来进行反向传播。这两者是类似的。在之前Variable的学习中，我们知道进行一次运算后，输出的Variable对应的creator就是其运行的计算，如y = relu(x), y.creator，就