torch.autograd.Function

import torch
from torch.autograd import Variable

class MyReLU(torch.autograd.Function):

    def forward(self, input_):
        # 在forward中，需要定义MyReLU这个运算的forward计算过程
        # 同时可以保存任何在后向传播中需要使用的变量值
        self.save_for_backward(input_)         # 将输入保存起来，在backward时使用
        output = input_.clamp(min=0)           # relu就是截断负数，让所有负数等于0
        return output

    def backward(self, grad_output):
        # 根据BP算法的推导（链式法则），dloss / dx = (dloss / doutput) * (doutput / dx)
        # dloss / doutput就是输入的参数grad_output、
        # 因此只需求relu的导数，在乘以grad_outpu    
        input_, = self.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0       # 上述计算的结果就是左式。即ReLU在反向传播中可以看做一个通道选择函数，所有未达到阈值（激活值<0）的单元的梯度都为0
        return grad_input

from torch.autograd import Variable
input_ = Variable(torch.randn(1))
relu = MyReLU()
output_ = relu(input_)

# 这个relu对象，就是output_.creator，即这个relu对象将output与input连接起来，形成一个计算图
print relu
print output_.creator


Output：
<__main__.MyReLU object at 0x7fd0b2d08b30>
<__main__.MyReLU object at 0x7fd0b2d08b30>

def relu(input_):
    # MyReLU()是创建一个MyReLU对象，
    # Function类利用了Python __call__操作，使得可以直接使用对象调用__call__制定的方法
    # __call__指定的方法是forward，因此下面这句MyReLU（）（input_）相当于
    # return MyReLU().forward(input_)
    return MyReLU()(input_)
input_ = Variable(torch.linspace(-3, 3, steps=5))
print input_
print relu(input_) 


Output：
Variable containing:
-3.0000
-1.5000
 0.0000
 1.5000
 3.0000
[torch.FloatTensor of size 5]

Variable containing:
 0.0000
 0.0000
 0.0000
 1.5000
 3.0000
[torch.FloatTensor of size 5]