自定义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] =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_()