几个例子帮你梳理PyTorch知识点（张量、autograd）

因为我最近想学Pytorch lightning，重构一下之前的代码，所以回来梳理一下Pytorch的语法，好进行下一步学习，所以从头重新回顾一下Pytorch。这个文章是通过几个简单例子帮大家回顾一下Pytorch一些重点基础概念。

Pytorch有两个重要的特征：

• 使用n维张量进行运算，可以使用GPU加速计算。
• 使用自动微分构建、训练神经网络

从$sin(x)$开始

Tensor和Numpy的用法差不多，但是Tensor可是使用进行加速计算，这比CPU计算要快50倍甚至更多。

更多区别可以看这个：PyTorch的Tensor这么简单，你还用不明白吗？ - 掘金 (juejin.cn)

我们知道，在基础的回归问题中：给定一些数据，符合一定的分布，我们要建立一个神经网络去拟合这个分布，神经网络学习出来的表达式就作为我们数据分布的表达式。

这里我们用一个三次多项式$y=a+bx+cx^2+dx^3$来拟合$sin(x)$，训练过程使用随机梯度下降进行训练，通过计算最小化预测值和真实值之间的欧氏距离来拟合随机数据。

代码解释见代码中的注释部分。

import torch
import math

dtype = torch.float
device = torch.device("cpu")


# 下边这行代码可以用也可以不用，注释掉就是在CPU上运行
# device = torch.device("cuda:0") 

# 创建输入输出数据，这里是x和y代表[-π，π]之间的sin(x)的值
x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
y = torch.sin(x)

#随机初始化权重
a = torch.randn((), device=device, dtype=dtype)
b = torch.randn((), device=device, dtype=dtype)
c = torch.randn((), device=device, dtype=dtype)
d = torch.randn((), device=device, dtype=dtype)

learning_rate = 1e-6
for t in range(2000):
    # 前向过程，计算y的预测值
    y_pred = a + b * x + c * x ** 2 + d * x ** 3

    # 计算预测值和真实值的loss
    loss = (y_pred - y).pow(2).sum().item()
    if t % 100 == 99:
        print(t, loss)

    # 反向过程计算 a, b, c, d 关于 loss 的梯度
    grad_y_pred = 2.0 * (y_pred - y)
    grad_a = grad_y_pred.sum()
    grad_b = (grad_y_pred * x).sum()
    grad_c = (grad_y_pred * x ** 2).sum()
    grad_d = (grad_y_pred * x ** 3).sum()

    # 使用梯度下降更新参数
    a -= learning_rate * grad_a
    b -= learning_rate * grad_b
    c -= learning_rate * grad_c
    d -= learning_rate * grad_d


print(f'Result: y = {a.item()} + {b.item()} x + {c.item()} x^2 + {d.item()} x^3')

结果如下：

如果换成勒让德多项式呢？

之前讲过深扒torch.autograd原理 - 掘金 (juejin.cn)

这里我们再浅浅简述一下autograd。

在上边的例子里，我们是手动实现了神经网络的前向和反向传播过程，因为这只是一个简单的两层网络，所以现实来也不是很困难，但是放对于一些大的复杂网络，要手动实现整个前向和反向过程就是非常困难的事情了。

现在我们可以使用pytorch提供autograd包去自动求导，自动计算神经网络的反向过程。当我们使用autograd的时候，神经网络前向过程就是定义一个计算图，计算图上的节点都是张量，边是从输入到输出的计算函数。用计算图进行反向传播可以轻松计算梯度。

虽然听起来很复杂，但是用起来很简单。每个张量都代表计算图上的一个节点，如果x是一个张量，并且你设置好了x.requires_grad=True，那x.grad就是另一个存储x关于某些标量的梯度的张量。

然后我们继续使用三次多项式来拟合我们的sin(x)，但是现在我们就可以不用手动实现反向传播的过程了。

import torch
import math
 
dtype = torch.float
device = torch.device("cpu")

# 下边这行代码可以用也可以不用，注释掉就是在CPU上运行
# device = torch.device("cuda:0") 
 
# 创建输入输出数据，这里是x和y代表[-π，π]之间的sin(x)的值
x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
y = torch.sin(x)
 
# 随机初始化权重
# 注意这里我们设置了requires_grad=True，让autograd自动跟踪计算图的梯度计算
a = torch.randn((), device=device, dtype=dtype, requires_grad=True)
b = torch.randn((), device=device, dtype=dtype, requires_grad=True)
c = torch.randn((), device=device, dtype=dtype, requires_grad=True)
d = torch.randn((), device=device, dtype=dtype, requires_grad=True)
 
learning_rate = 1e-6
for t in range(2000):
    
    # 前向过程，计算y的预测值
    y_pred = a + b * x + c * x ** 2 + d * x ** 3

    # 计算预测值和真实值的loss
    loss = (y_pred - y).pow(2).sum()
    if t % 100 == 99:
        print(t, loss.item())
 
    # 使用autograd计算反向过程，调用之后会计算所有设置了requires_grad=True的张量的梯度
    # 调用之后 a.grad, b.grad. c.grad  d.grad 会存储 abcd关于loss的梯度
    loss.backward()
 
    # 使用梯度下降更新参数
    # 因为权重设置了requires_grad=True，但是在梯度更新这里我们不需要跟踪梯度，所以加上with torch.no_grad()
    with torch.no_grad():
        a -= learning_rate * a.grad
        b -= learning_rate * b.grad
        c -= learning_rate * c.grad
        d -= learning_rate * d.grad
 
        # 更新之后将气度清零，以便下一轮运算，不清零的话它会一直累计
        a.grad = None
        b.grad = None
        c.grad = None
        d.grad = None
 
print(f'Result: y = {a.item()} + {b.item()} x + {c.item()} x^2 + {d.item()} x^3')

在pytorch这种autograd的情况下，每个基础的的autograd操作只是两个作用于张量的方法。

• forward：由输入张量计算输出张量
• backward：接收输出张量相对于某个标量值的梯度，并计算输入张量相对于相同标量值的梯度。

在pytorch中我们可以定义我们自己的autograd操作，你只需要实现一个torch.autograd.Function的子类，写好forward和backward函数即可。构造好新的autograd之后我们就可以像调用其他函数一样调用它，将输入张量传递进去即可。

比如我们不用$y=a+bx+cx^2+dx^3$了，改成一个三次勒让德多项式（Legendre polynomial），形式为$y = a+bP_3(c+dx)$，其中$P_3(x) = \frac 1 2 (5x^3-3x)$。

import torch
import math
 
class LegendrePolynomial3(torch.autograd.Function):
    def forward(ctx, input):
        
        # 在前向过程中我们接受一个输入张量，并返回一个输出张量
        # ctx是一个上下文对象，用于存储反向过程的内容
        # 你可以用save_for_backward方法缓存任意在反向计算过程中要用的对象。
        
        ctx.save_for_backward(input)
        return 0.5 * (5 * input ** 3 - 3 * input)

    def backward(ctx, grad_output):
        # 在反向过程中，我们接受一个张量包含了损失关于输出的梯度，我们需要计算损失关于输入的梯度。
        input, = ctx.saved_tensors
        return grad_output * 1.5 * (5 * input ** 2 - 1)
 
        
dtype = torch.float
device = torch.device("cpu")

# 下边这行代码可以用也可以不用，注释掉就是在CPU上运行
# device = torch.device("cuda:0") 

# 创建输入输出数据，这里是x和y代表[-π，π]之间的sin(x)的值
x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
y = torch.sin(x)

# 随机初始化权重
# 注意这里我们设置了requires_grad=True，让autograd自动跟踪计算图的梯度计算
a = torch.full((), 0.0, device=device, dtype=dtype, requires_grad=True)
b = torch.full((), -1.0, device=device, dtype=dtype, requires_grad=True)
c = torch.full((), 0.0, device=device, dtype=dtype, requires_grad=True)
d = torch.full((), 0.3, device=device, dtype=dtype, requires_grad=True)
 
learning_rate = 5e-6
for t in range(2000):
    # 我们给我们自定义的autograd起个名叫P3，然后用Function.apply方法调用
    P3 = LegendrePolynomial3.apply
 
    # 前向过程计算y，用的是我们自定义的P3的autograd
    y_pred = a + b * P3(c + d * x)
 
    # 计算并输出loss
    loss = (y_pred - y).pow(2).sum()
    if t % 100 == 99:
        print(t, loss.item())
 
    # 使用autograd计算反向过程
    loss.backward()
 
    # 使用梯度下降更新权重
    with torch.no_grad():
        a -= learning_rate * a.grad
        b -= learning_rate * b.grad
        c -= learning_rate * c.grad
        d -= learning_rate * d.grad
 
        # 在下一轮更新之前将梯度清零，否则会一直累计
        a.grad = None
        b.grad = None
        c.grad = None
        d.grad = None
 
print(f'Result: y = {a.item()} + {b.item()} * P3({c.item()} + {d.item()} x)')