[官网] 优化模型参数

现在有模型和数据，可以通过优化模型参数去训练，验证和测试模型。训练模型是一个迭代的过程，在一次迭代中，模型会关于输出做出猜测，利用猜测计算差距（损失），收集模型参数的损失，然后使用梯度下降优化这些参数。

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

training_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 logits

model = NeuralNetwork()

超参数

超参数是来控制模型优化过程的可调整参数，不同的超参数值能影响模型训练和收敛比率。

定义了以下超参数来训练：

• Number of Epochs：迭代整个数据集的次数
• Batch Size：在参数更新前，前向传播神经网络的数据样本的数据量。
• Learning Rate：每次batch/epoch，更新模型参数多少。

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

优化循环

一旦设置了超参数，就可以训练和优化模型一次优化循环。每一次优化循环的迭代称作为epoch。

每次epoch由两个主要部分：

• 训练循环：迭代训练数据集并且收敛优化参数。
• 验证/测试循环：迭代测试数据集去测试模型性能是否在提高。

损失函数

损失函数是判定获得结果与目标值差异的程度。为了计算损失，使用提供的数据样本输入，做出预测，与真实的数据标签值进行比较。

通用损失函数：

• nn.MESLoss(Mean Square Error)：均方误差，处理回归任务。
• nn.NLLLoss(Negative Log Likeihood)：负对数似然函数，处理分类。
• nn.CrossEntropyLoss：交叉熵函数，结合nn.LogSoftMax和nn.NLLLoss。

这里用的是nn.CrossEntropyLoss

# Initialize the loss function
loss_fn = nn.CrossEntropyLoss()

优化器

优化是在一次训练循环，调整模型参数，减少模型误差的过程。优化算法定义了这个过程是怎么执行的（例子中使用的是Stochastic Gradient Descent 随机梯度下降算法）。所有的优化逻辑都是包装在优化器对象中。

通过注册需要训练的模型参数，和传递学习率超参数，来初始化优化器。

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

在训练循环中，优化发生了三个步骤：

• 调用optimizer.zero_grad()去重置模型参数的梯度，梯度默认是累加的，为了避免双重计数，每次迭代明确地置0。
• 调用loss.backward()反向传播预测损失，pytorch存储每个参数损失的梯度。
• 调用optimizer.step()去通过在反向传递过程中收集到的梯度调整参数。

完整实现

定义train_loop循环优化代码，定义test_loop通过测试数据来评估模型性能。

def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    # Set the model to training mode - important for batch normalization and dropout layers
    # Unnecessary in this situation but added for best practices
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        # Compute prediction and loss
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * batch_size + len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")


def test_loop(dataloader, model, loss_fn):
    # Set the model to evaluation mode - important for batch normalization and dropout layers
    # Unnecessary in this situation but added for best practices
    model.eval()
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    # Evaluating the model with torch.no_grad() ensures that no gradients are computed during test mode
    # also serves to reduce unnecessary gradient computations and memory usage for tensors with requires_grad=True
    with 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_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

初始化损失函数和优化器，传递它到train_loop和test_loop。提高epochs的数量来追踪模型提高性能。

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

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!")