pytorch基础五（定义自动求导函数）

本人学习pytorch主要参考官方文档和 莫烦Python中的pytorch视频教程。
后文主要是对pytorch官网的文档的总结。
代码来自pytorch官网

import torch
# 通过继承torch.autograd.Function类，并实现forward 和 backward函数
class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        """
        在forward函数中，接收包含输入的Tensor并返回包含输出的Tensor。
        ctx是环境变量，用于提供反向传播是需要的信息。可通过ctx.save_for_backward方法缓存数据。
        """
        ctx.save_for_backward(input)
        return input.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        """
        在backward函数中，接收包含了损失梯度的Tensor，
        我们需要根据输入计算损失的梯度。
        """
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input

dtype = torch.float
device = torch.device("cpu")
N, D_in, H, D_out = 64, 1000, 100, 10
x = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)
w1 = torch.randn(D_in, H, device=device, dtype=dtype, requires_grad=True)
w2 = torch.randn(H, D_out, device=device, dtype=dtype, requires_grad=True)
learning_rate = 1e-6
for t in range(500):
    relu = MyReLU.apply
    y_pred = relu(x.mm(w1)).mm(w2)
    loss = (y_pred - y).pow(2).sum()
    print(t, loss.item())
    loss.backward()
    with torch.no_grad():
        w1 -= learning_rate * w1.grad
        w2 -= learning_rate * w2.grad
        w1.grad.zero_()
        w2.grad.zero_()