当前位置：首页 > news >正文

Pytorch-06 如何开启模型训练（模型训练工作流梳理）

news 2025/8/11 5:48:53

开启训练-优化模型参数

经过我们前面的努力，现在我们已经准备好了网络模型和数据。

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensortraining_data = datasets.FashionMNIST(root="data",train=True,download=True,transform=ToTensor()
)test_data = datasets.FashionMNIST(root="data",train=False,download=True,transform=ToTensor()
)train_dataloader = DataLoader(training_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)class NeuralNetwork(nn.Module):def __init__(self):super().__init__()self.flatten = nn.Flatten()self.linear_relu_stack = nn.Sequential(nn.Linear(28*28, 512),nn.ReLU(),nn.Linear(512, 512),nn.ReLU(),nn.Linear(512, 10),)def forward(self, x):x = self.flatten(x)logits = self.linear_relu_stack(x)return logitsmodel = NeuralNetwork()

现在是时候开启我们的训练过程来优化带学习的参数了。

训练模型是一个迭代的过程，每次迭代中，模型都会猜一个输出，计算其对于真实标签的loss，然后根据loss反向传播更新梯度到需要学习的参数上，然后对它们使用 梯度下降算法 去优化它们。你可以看看3B1B的视频去更详细地理解这个过程：backpropagation from
3Blue1Brown.

接下来我们将列出开启训练之前要做的必要设置。

1. 训练超参数的设置

超参数是一种你可以调整的，可以让你控制模型优化过程的参数，以下是三个最经典的，不管你是什么训练都跑不掉设置的超参数：

learning_rate: 学习率，每次迭代的时候，有了更新方向之后（梯度反方向），要往这个方向走多远
batch_size：一次迭代中一个minibatch包含多少个样本
epochs: 整个数据集你要跑几轮

learning_rate = 1e-3
batch_size = 64
epochs = 5

2. 定义优化循环

一旦我们设置好了训练超参数，我们就可以开始用优化循环来训练我们的模型了。优化循环的一次迭代被成为epoch，也就是跑一整轮数据集就是一次优化循环的迭代。

每一个epoch包括两个主要部分

Train Loop：训练循环，负责在训练集上进行迭代并优化学习参数
Validation Loop / Test Loop: 验证/测试循环在验证集或者测试机上遍历推理，获得最终的模型效果数据

我们先过一下train loop里面会用到的关键概念

2.1 Loss Function 损失函数

当我们给还没训练的模型一些训练的数据的时候，它不太会给出一个比较符合真实情况的答案，那么这个答案有多“离谱”呢？这就需要损失函数来测量了。损失函数接受真实标签和预测标签，用一个预先准备好的函数来计算它两之间的差异。

比较常见的损失函数有

nn.MSELoss 平方差损失（用于回归）
nn.NLLLose 负对数似然损失（用于分类）

这里我们用的nn.CrossEntropyLoss则是nn.LogSoftMax以及nn.NLLLoss的结合体:

loss_fn = nn.CrossEntropyLoss()

其公式如下所示：
在这里插入图片描述
我比较喜欢这样来直观理解为什么交叉熵损失是这样的：

2.2 Optimizer 优化器

优化器决定了每个训练步会如何根据梯度来进行参数的调整。不同的优化器会有不同的优化算法封装所有的优化逻辑在自己的实例中。这里我们使用的是SGD随机梯度下降：

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

其包含的优化算法流程如下：在这里插入图片描述
在训练循环中，优化过程被分为三步来走：

调用optimizer.zero_grad()：梯度清零。将之前一次迭代造成的梯度清空，为下一次的迭代所算出的梯度腾出空间。
调用loss.backward()：反向传播，把所有需要学习的参数在这次迭代的梯度算出来
调用optimizer.step()：每个要学习的参数有了一度之后，使用优化器步进，执行每个参数的更新逻辑以调整参数。

还有很多其他的优化器，例如Adam，RMSProp等等，可以在这里看到： different
optimizers

3. 模型训练完整实现

我们一般定义两个loop

一个是train_loop，也就是训练循环，其中要执行前向预测和更新参数的逻辑。
另一个是test_loop，也就是测试循环，其中只需要执行前向预测即可。

def train_loop(dataloader, model, loss_fn, optimizer):size = len(dataloader.dataset)model.train()for batch, (X, Y) in enumerate(dataloader):pred = model(X)loss = loss_fn(pred, y)loss.backward()optimizer.step()optimizer.zero_grad()if batch % 100 == 0:loss, current = loss.item(), batch * batch_size + len(X) # current是当前到第几个样本了print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")， # loss右对齐，至少7个字符，浮点数，current/size右对齐，至少5个字符，浮点数def test_loop(dataloader, model, loss_fn):size = len(dataloader.dataset)model.eval() # 启动评估模式num_batches = len(dataloader)test_loss, correct = 0, 0with torch.no_grad():for X, y in dataloader:pred = model(X)test_loss += loss_fn(pred, y).item()correct += (pred.argmax(1) == y).type(torch.float).sum().item()test_loss /= num_batchescorrect /= sizeprint(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

定义完这两个循环之后，就可以开启训练了：

epochs = 10
for t in range(epochs):print(f"Epoch {t+1}\n--------------")train_loop(train_dataloader, model, loss_fn, optimizer)test_loop(test_dataloader, model, loss_fn)
print("Done!")