引言
本文的目的是在前文基于numpy演练可视化梯度下降的代码基础上，使用pytorch来实现一个功能齐全的线性回归训练模型。
为什么仍然使用线性回归模型？
- 线性回归模型简单，它能让我们聚集在pytorch是如何工作的，而不是模型内部的某个复杂结构或算法。
- 与前面的[基于numpy的线性回归模型]作对比，pytorch如何让代码量更少，更易于理解。
1. 数据准备
先导入整体要用到的包。
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
1.1 数据生成
数据生成和之前一样，这里不再赘述。
true_w = 2
true_b = 1
N = 100
np.random.seed(42)
x = np.random.rand(N, 1)
eplison = 0.1 * np.random.randn(N, 1)
y = true_w * x + true_b + eplison
x.shape, y.shape, eplison.shape
((100, 1), (100, 1), (100, 1))
数据集拆分也没什么变化。
idx = np.arange(N)
np.random.shuffle(idx)
ratio = int(0.8 * N)
x_train, x_test = x[idx[:ratio]], x[idx[ratio:]]
y_train, y_test = y[idx[:ratio]], y[idx[ratio:]]
x_train.shape, y_train.shape, x_test.shape, y_test.shape
((80, 1), (80, 1), (20, 1), (20, 1))
1.2 数据转换
主要是两方面的转换：
- 设备：之前使用numpy时不用关心设备（numpy只支持cpu），现在使用pytorch时，需要指定设备。
- 数据类型：矩阵转换为张量tensor。
device = 'cuda' if torch.cuda.is_available() else 'cpu'
x_train_tensor = torch.from_numpy(x_train).float().to(device)
y_train_tensor = torch.from_numpy(y_train).float().to(device)
x_test_tensor = torch.from_numpy(x_test).float().to(device)
y_test_tensor = torch.from_numpy(y_test).float().to(device)
x_train.shape, x_train_tensor.shape, y_train_tensor.shape
((80, 1), torch.Size([80, 1]), torch.Size([80, 1]))
上面将整个数据集发送到device的做法，在实际中可能存在隐患，因为数据集可能很大，如果直接发送到device上，可能会占用宝贵的显存空间。
1.3 定义数据集
在pytorch中，数据集是torch.utils.data.Dataset的子类，可以把它理解成一个元组列表，每个元组对应一个点，包含特征x、标签y。使用Dataset类需要重写几个方法：
- init：初始化，它可以接收数据文件的路径，也可以直接接收两个张量x和y，分别表示特征和标签。
- getitem(index)：通过索引下标对数据集进行访问，可以访问单个数据点、数据切片、或者按需加载，但有一点要求是必须返回包含特征和标签的元组。
- len：返回整个数据集的大小。
使用Dataset的好处是：可以不用一次性加载整个数据集，而是每当调用__getitem__方法时，按需加载。
并且重写了__len__和__getitem__方法。
from torch.utils.data import Dataset, DataLoader
class MyDataset(Dataset):
def __init__(self, x_data, y_data):
self.x = x_data
self.y = y_data
def __getitem__(self, index):
return (self.x[index], self.y[index])
def __len__(self):
return len(self.x)
train_dataset = MyDataset(x_train_tensor, y_train_tensor)
test_dataset = MyDataset(x_test_tensor, y_test_tensor)
train_dataset[0], train_dataset[:5]
((tensor([0.7713]), tensor([2.4745])),
(tensor([[0.7713],
[0.0636],
[0.8631],
[0.0254],
[0.7320]]),
tensor([[2.4745],
[1.1928],
[2.9128],
[1.0785],
[2.4732]])))
1.4 定义小批量数据加载器
前面验证了小批量梯度下降，在同样的数据量下，比批量梯度更容易收敛，比随机梯度下降计算量更少，效果更稳定。小批量梯度下降最主要的工作就是选择每次训练使用多少数据量，以及使用哪部分数据集，pytorch的DataLoader类可以很方便完成这项工作，只需要为它传3个参数：
- dataset: 数据集
- batch_size: 小批量大小
- shuffle: 是否打乱数据，默认是False
注：绝大多数情况下，我们都应该把shuffle设为True，以提高梯度下降性能。但是验证集和测试集其实没必要打乱，因为它们并不参与梯度计算。
注：小批量大小，通常使用2的冥，如8、16、32、64等，这样能进行内存对齐，因为CPU/GPU的内存架构通常都是按照2的冥来分配内存。
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=8, shuffle=True)
# iter函数用于获取一个迭代器，访问迭代器可以逐个获取每个数据批次
# next函数用于从迭代器中获取下一个数据批次
x, y = next(iter(train_loader))
x, y
(tensor([[0.2809],
[0.8631],
[0.3110],
[0.9507],
[0.0740],
[0.2912],
[0.6233],
[0.1220]]),
tensor([[1.5846],
[2.9128],
[1.5245],
[2.8715],
[1.1713],
[1.4361],
[2.2940],
[1.2406]]))
2. 模型配置
2.1 创建参数
训练数据和参数/权重都是tensor， 但两者最大的区别在于：参数/权重是有梯度的，可以更新参数值，而requires_grad=True就是告诉pytorch，这是一个可学习的参数，需要梯度计算。
pytorch中也有与numpy中相似的api来设置随机数种子，并随机初始化参数，唯一不同的是，我们需要通过device参数将创建的参数分配到指定设备上。
torch.manual_seed(42)
w = torch.randn(1, requires_grad=True, dtype=torch.float, device=device)
b = torch.randn(1, requires_grad=True, dtype=torch.float, device=device)
w, b
(tensor([0.3367], requires_grad=True), tensor([0.1288], requires_grad=True))
2.2 定义模型
在pytorch中，模型由继承自nn.Module的类来表示，模型类最基本的两个方法是：
- init：定义构成模型的组成部分，包括定义参数，以及嵌套其它模型。
- forward(x)：定义模型的前向传播过程，即实际的计算操作，给定输入x的情况下，输出一个预测。
class LinearModel(nn.Module):
def __init__(self):
super().__init__()
self.w = nn.Parameter(torch.randn(1, requires_grad=True, dtype=torch.float, device=device))
self.b = nn.Parameter(torch.randn(1, requires_grad=True, dtype=torch.float, device=device))
def forward(self, x):
return self.w * x + self.b
torch.manual_seed(42)
model = LinearModel().to(device)
model.w, model.b
# model.state_dict
(Parameter containing:
tensor([0.3367], requires_grad=True),
Parameter containing:
tensor([0.1288], requires_grad=True))
注：在使用模型进行预测时，应该调用model(x)而不是model.forward(x)，原因是对整个模型的调用涉及到额外的步骤
注意到，在__init__方法中使用Parameter类包装了参数b和w，这样做的目的是可以使用模型的parameters()来检索所有模型参数，甚至包括嵌套模型的参数，而不用自己构建参数列表，这在复杂模型中非常有用。
list(model.parameters())
[Parameter containing:
tensor([0.3367], requires_grad=True),
Parameter containing:
tensor([0.1288], requires_grad=True)]
此外，还可以使用模型的state_dict方法来获取所有参数的当前值。
model.state_dict()
OrderedDict([('w', tensor([0.3367])), ('b', tensor([0.1288]))])
2.3 定义损失函数
pytorch中的nn.MSELoss()函数，可以返回计算均方误差的函数lossfn。MSELoss是一个高阶函数，支持通过reduction参数来指定损失的计算方式，默认值为mean表示计算均方差。
lossfn = nn.MSELoss()
yhat = model(x_train_tensor)
loss = lossfn(yhat, y_train_tensor)
loss
tensor(3.0332, grad_fn=<MseLossBackward0>)
2.4 自动梯度
在pytorch中，梯度是不需要手动计算的，它通过反向传播函数backward()就能自动计算所有（需要梯度的）张量的梯度。
loss.backward()
model.w.grad, model.b.grad
(tensor([-1.8803]), tensor([-3.3318]))
它代替了numpy版本中手动计算梯度的代码：
b_grad = 2 * error.mean()
w_grad = 2 * (x_train * error).mean()
单从这个例子可能还不足以看出反向传播函数的好处，试想如果是包含几十个层次的神经网络，每一层的参数都去手动推导公式并计算梯度，将是多么复杂的一件事。但现在一句代码
loss.backward()就将所有参数的梯度都计算好了。
loss.backward()之所以能做到这一切，是基于梯度计算的链式法则，详细请参考：动手学深度学习-求导
这里需要说明的一点是，pytorch中的梯度是会累加的，如果将上面计算梯度和反向传播的代码再运行一次，会发现梯度变成原来的2倍。
yhat = model(x_train_tensor)
loss = lossfn(yhat, y_train_tensor)
loss.backward()
model.w.grad, model.b.grad
(tensor([-3.7606]), tensor([-6.6637]))
这会带来一个问题，在使用小批量梯度下降进行训练时，第2个小批量的梯度会在第1个小批量梯度的基础上累加。但是，我们希望每个小批量上的梯度都应该是基于当前损失独立计算的，不应该使用累积梯度。
因此，我们每轮使用梯度更新参数后，需要将梯度清零。
model.w.grad.zero_()
model.b.grad.zero_()
model.w.grad, model.b.grad
(tensor([0.]), tensor([0.]))
2.5 动态计算图
动态计算图是以可视化的方式来展示模型结构和计算过程，主要用到两个软件包：torchviz和graphviz。torchviz安装比较简单，graphviz的安装参考：Mac下安装Graphviz实用教程。
torchviz的使用比较简单，只需要调用make_dot函数，并传入预测值yhat即可。
from torchviz import make_dot
make_dot(yhat)
- 蓝色框对应于参数w和b，就是需要计算梯度的张量。
- 灰色框（MulBackward0和AddBackward0）是计算梯度时，需要临时保存的中间结果。
- 绿色框（80,1）是梯度计算起点的张量，也就是调用backward函数的loss,反向传播是自下而上的。
为什么没有一个数据框（x）呢？
原因在于，x是输入数据，不需要计算梯度。计算图只显示涉及梯度计算的张量及其依赖的计算关系。
如果我们将参数b设为不需要梯度，那么b所在的计算分支将从计算图中消失。
model.b.requires_grad_(False)
yhat = model(x_train_tensor)
make_dot(yhat)
测试过后，记得将参数b的requires_grad恢复为True。
2.6 定义参数优化器
就和梯度计算一样，当涉及到复杂模型中有很多参数时，手动更新参数将不现实。pytorch中提供了optim模块，里面有很多优化器可以来更新参数，我们这里就使用随机梯度下降SGD优化器。
- 只要指定参数和学习率，再调用step方法就可以自动更新参数。
- 调用zero_grad方法可以将所有参数的梯度置零，不再需要逐个调用每个参数梯度的
_zero方法。
optimizer = optim.SGD([model.w, model.b], lr=0.2)
optimizer.step()
optimizer.zero_grad()
model.w, model.b, model.w.grad, model.b.grad
(Parameter containing:
tensor([0.3367], requires_grad=True),
Parameter containing:
tensor([0.1288]),
None,
None)
3 训练
3.1 构建训练步骤
在训练阶段，其实就是固定的4个步骤：
- 计算模型预测值
- 计算损失值
- 计算梯度
- 更新参数
这4个步骤会在不同的迭代数据集上反复执行，所以我们有必要封装一个函数，方便我们重复使用这段逻辑。
def build_train_step(model, loss_fn, optimizer):
def train_step(x, y):
# 设置模型为训练模式，
model.train()
# 模型预测——前向传递
yhat = model(x)
# 计算损失
loss = loss_fn(yhat, y)
# 反向传播计算梯度
loss.backward()
# 使用梯度和学习率更新参数
optimizer.step()
optimizer.zero_grad()
# 返回损失
return loss.item()
return train_step
注：模型、损失函数、优化器是会改变的，所以这几个将作为参数来构建训练步骤train_step, 返回的训练步骤train_step接受特征x和标签y作为参数来完成一轮训练。
3.2 构建验证步骤
验证模型的目的，是为了观察模型对从未见过数据进行预测时的错误程度。验证步骤和训练步骤非常相似，但是不需要梯度下降。
- 使用模型来计算预测
- 使用损失函数来计算损失
还有一重要区别：必须调用eval方法将模型设置评估模式。
在训练模式时，模型会自动执行一些操作（如dropout丢弃）来减少过拟合，这种操作会破坏评估，所以在评估模式时需要关掉。
def build_evaluate_step(model, loss_fn):
def evaluate_step(x, y):
# 设置模型为评估模式
model.eval()
# 计算模型的预测输出，前向传递
y_hat = model(x)
# 计算损失
loss = loss_fn(y_hat, y)
# 返回损失值
return loss.item()
return evaluate_step
3.3 训练循环
当封装了训练步骤后，训练循环主要就做三件事：
- 迭代数据
- 执行一个训练步骤
- 跟踪训练损失
而对于模型验证来说，采用边训练边验证损失更容易发现一些过拟合的问题，所以验证损失也作为训练的一部分，包括执行步骤也是上面三步，唯一的区别在于把它包裹在了torch.no_grad()中。
torch.no_grad()相当于一个上下文管理器，它能禁用任何训练阶段的操作，避免在验证阶段误触发会对模型产生影响的梯度计算，同时也节省时间和计算资源。
epoch = 100
losses, test_losses = [], []
w_history, b_history = [model.w.item()], [model.b.item()]
train_step = build_train_step(model, lossfn, optimizer)
evaluate_step = build_evaluate_step(model, lossfn)
for i in range(epoch):
# 迭代下一个小批量数据集
x, y = next(iter(train_loader))
# 执行一个训练步骤
loss_val = train_step(x.to(device), y.to(device))
# 记录训练损失
losses.append(loss_val)
# 记录模型参数
w_history.append(model.w.item())
b_history.append(model.b.item())
# 验证时，不更新模型参数
with torch.no_grad():
x, y = next(iter(test_loader))
test_loss = evaluate_step(x.to(device), y.to(device))
test_losses.append(test_loss)
model.state_dict(), optimizer.state_dict(), losses, test_losses,
(OrderedDict([('w', tensor([1.9258])), ('b', tensor([1.0505]))]),
{'state': {0: {'momentum_buffer': None}, 1: {'momentum_buffer': None}},
'param_groups': [{'lr': 0.2,
'momentum': 0,
'dampening': 0,
'weight_decay': 0,
'nesterov': False,
'maximize': False,
'foreach': None,
'differentiable': False,
'params': [0, 1]}]},
[2.3343870639801025,
1.4687697887420654,
0.15246182680130005,
……
0.013301181606948376,
0.009759355336427689],
[1.1239256858825684,
0.0859474241733551,
0.14135879278182983,
0.011405732482671738,
……
0.007371845189481974])
def show_losses(losses, test_losses, w_history, b_history):
fig, ax = plt.subplots(1, 2, figsize=(12, 6))
epoches = range(1, len(losses) + 1)
ax[0].plot(epoches, losses, label='train losses')
ax[0].plot(epoches, test_losses, label='test losses')
ax[0].set_xlabel('epoch')
ax[0].set_ylabel('loss')
ax[0].set_yscale('log') #
ax[0].set_title('loss descent path')
ax[0].legend()
ax[1].plot(b_history, w_history, c='b', marker='.', linewidth=0.5, linestyle='--')
ax[1].set_xlabel('b')
ax[1].set_ylabel('w')
ax[1].set_title('parameters fitting path')
ax[1].annotate(f'Random start({b_history[0]:0.4f}, {w_history[0]:0.4f})', xy=(b_history[0]+0.1, w_history[0]), fontsize=10, color='k')
ax[1].plot(true_b, true_w, c='k', marker='o')
ax[1].annotate(f'True values({true_b, true_w})', xy=(true_b+0.1, true_w-0.01), fontsize=10, color='g')
plt.tight_layout()
plt.show()
show_losses(losses, test_losses, w_history, b_history)
注：可以看到损失的y轴采用的是对数，原因在于损失的数值不是线性均匀分布，最开始的几轮训练存在少数损失较大的值，后面更多轮训练的损失都集中在了某个小区域内，如果采用线性分布会使得图形的显示效果不佳。