PyTorch实战第一课：训练三阶多项式

任务

对于函数$y = sin(x)$，在$-\pi$到$\pi$的区间上，我们希望建立一个模型去学习（拟合）它，模型为：

\[\hat{y}=ax^3+bx^2+cx+d\]

其中，参数$a,b,c,d$需要模型自己学习。

准备数据

在已经知道目标函数的情况下，数据准备很简单，直接按照原函数生成：

x = torch.linspace(-math.pi, math.pi, 2000)
# x是一个tensor，包含2000个元素，均匀分布从-pi到pi
y = torch.sin(x)

定义网络

虽然这是一个简单的任务，但是这篇文章的目的是简单使用并熟悉PyTorch为我们写好的各种“轮子”。定义网络将使用torch.nn提供的类和方法：

class Polynomial3(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.a = torch.nn.Parameter(torch.randn(()))
        self.b = torch.nn.Parameter(torch.randn(()))
        self.c = torch.nn.Parameter(torch.randn(()))
        self.d = torch.nn.Parameter(torch.randn(()))

    def forward(self, x):
        return self.a + self.b * x + self.c * x ** 2 + self.d * x ** 3


model = Polynomial3()

torch.nn.Module：创建一个行为类似于函数的可调用对象，但也可以包含状态（例如神经网络层权重）。 它知道其中包含的 Parameter ，并且可以将其所有坡度归零，遍历它们以进行权重更新等。

torch.nn.Parameter：张量的包装器，用于告知 Module 具有在反向传播期间需要更新的权重。 仅更新具有require_grad属性集的张量

• __init__(self)函数负责初始化参数和定义要使到的层。
• forward(self, x)函数负责定义$x$进入网络的计算过程。可以看到这个例子中$\hat{y}=ax^3+bx^2+cx+d$
• 定义完类之后，初始化一个对象

定义损失函数和优化器

criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=1e-6)

损失函数torch.nn.MSELoss使用输入$x$和目标$y$之间的均方误差，reduction='sum'表示不取平均值只是求和。

\[loss(x,y)=\sum_{i=1}^N(x_i-y_i)^2\]

torch.optim.SGD使用随机梯度下降算法作为优化器。lr代表学习率。

官方文档对于随机梯度下降算法的解释：https://pytorch.org/docs/stable/generated/torch.optim.SGD.html#torch.optim.SGD

实现单次训练

我们已经准备好了神经网络的所有东西，现在开始训练吧。下面是单次训练的流程，如果多次训练，循环就行了。

y_pred = model(x)
loss = criterion(y_pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()

1. 首先计算出预测值y_pred
2. 通过损失函数计算损失
3. 将梯度初始化为零，因为一个batch的loss关于weight的导数是batch中所有sample的loss关于weight的导数的累加和
4. 反向传播
5. 更新参数

组装到一起

import torch
import math


class Polynomial3(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.a = torch.nn.Parameter(torch.randn(()))
        self.b = torch.nn.Parameter(torch.randn(()))
        self.c = torch.nn.Parameter(torch.randn(()))
        self.d = torch.nn.Parameter(torch.randn(()))

    def forward(self, x):
        return self.d + self.c * x + self.b * x ** 2 + self.a * x ** 3

    def string(self):
        return f'y = {self.d.item()} + {self.c.item()} x + {self.b.item()} x^2 + {self.a.item()} x^3'


x = torch.linspace(-math.pi, math.pi, 2000)
y = torch.sin(x)

model = Polynomial3()
criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.SGD(model.parameters(), lr=1e-6)

for t in range(2000):
    y_pred = model(x)
    loss = criterion(y_pred, y)
    if t % 100 == 0:
        print(t, loss.item())
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

print(f'Result: {model.string()}')

新加入的东西：

• 网络的string(self)方法，输出五个参数的信息
• for t in range(2000): 循环训练2000次
• 训练每隔100次输出$loss$的值

self.d.item()的解释：item()方法将一个tensor转为python中的数，只对包含一个元素的tensor有效

输出

运行上面代码的输出如下：

0 949606.8125
100 154.83355712890625
200 111.57437133789062
300 81.15692901611328
400 59.75954055786133
500 44.70195770263672
600 34.102272033691406
700 26.63823890686035
800 21.380661010742188
900 17.676284790039062
1000 15.065535545349121
1100 13.225081443786621
1200 11.927343368530273
1300 11.012075424194336
1400 10.36642074584961
1500 9.910867691040039
1600 9.589386940002441
1700 9.362482070922852
1800 9.202301025390625
1900 9.089205741882324
Result: y = -0.014516396448016167 + 0.854619026184082 x + 0.002504321513697505 x^2 + -0.09302857518196106 x^3

• 蓝色为$y = sin(x)$
• 红色为$\hat{y} = -0.014516396448016167 + 0.854619026184082 x + 0.002504321513697505 x^2 + -0.09302857518196106 x^3$