自定义autograd

import torch

class MyreLu(torch.autograd.Function):
    """
    我们可以通过定义torch.autograd的子类来实现我们自定义的autograd函数,
    并完成张量的正向反向传播
    """
    @staticmethod
    def forward(ctx, x ):
        """
        在正向传播的时候，forward函数会接受一个上下文对象和一个包含输入输出的张量；
        我们必须返回一个输出张量；
        并且我们可以使用上下文对象来缓存对象，以便在反向传播中使用
        """
        print(x.size())
        ctx.save_for_backward(x)
        '''
        save_for_backward() 只能存储 tensor, None 其余的都不可以存储
        save_for_backward() 只能存储 forward 的实例参数，和forward的返回值
        '''
        return x.clamp(min=0) #实现relu的操作
    @staticmethod
    def backward(ctx, grad_output):
        """
        在反向传播的过程中，我们会接收上下文对象ctx和一个张量
        其中包含了相对于正传播中产生的输出损失的梯度。
        我们可以从上下文对象中检索缓存的数据
        并且计算与正向传播的输入相关的损失梯度
        """
        x,= ctx.saved_tensors  #这里在x后面为啥加逗号，博主也不知道，目前没搞明白
        print(x.size())

        #ctx.save_for_backward保存反向传播需要用到的参数；
        #ctx.saved_tensors读取参数
        grad_x = grad_output.clone()
        print(grad_x.size())
        grad_x[x < 0] =0    #这个表达式操作博主也不是很明白
        return grad_x
'''这里只是用了一个输入做一个小演示，正常的话是需要input，weight，bias 这三个参数的。
grad_x[x < 0] =  这个操作博主认为应该是判断其requires_grad是否为True，从而判断是否进行梯度的运算
但是至于这个为啥这么写，博主也不是很清楚

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

N, D_in, H, D_out = 64, 1000, 100, 10

x = torch.randn(N, D_in, device=device)
y = torch.randn(N, D_out, device=device)

w1 = torch.randn(D_in, H, device=device,requires_grad=True)
w2 = torch.randn(H, D_out, device=device,requires_grad=True)

learning_rate = 1e-6

for t in range(500):

    y_pred = MyreLu.apply(x.mm(w1)).mm(w2)
    loss = (y_pred-y).pow(2).sum()
    print(t,loss)
    loss.backward()

    with torch.no_grad():
        w1 -= learning_rate*w1.grad
        w2 -= learning_rate*w2.grad

        #完成更新后清零梯度

        w1.grad.zero_()
        w2.grad.zero_()