用 PyTorch 在 MNIST 数据集上进行逻辑回归

人工智能

2023-09-05 12:53:34

逻辑回归是一种强大且易于实现的机器学习算法，广泛应用于各种分类任务中。在这个教程中，我们将使用流行的 PyTorch 深度学习库在经典的 MNIST 手写数字数据集上构建一个逻辑回归模型。

逻辑回归概述

逻辑回归是一个二元分类算法，它预测给定输入属于两个可能类别之一的概率。与线性回归不同，逻辑回归使用非线性 sigmoid 函数将输入数据映射到概率空间。该函数将实数值输入转换为 0 到 1 之间的概率。

PyTorch 中的逻辑回归

PyTorch 是一个强大的 Python 库，用于构建和训练神经网络。它提供了高效且易于使用的工具，可以轻松实现各种机器学习模型，包括逻辑回归。

模型定义

要使用 PyTorch 定义逻辑回归模型，我们可以使用 torch.nn.Linear 和 torch.nn.Sigmoid 模块。Linear 模块表示线性变换，而 Sigmoid 模块应用 sigmoid 函数。

import torch
import torch.nn as nn
import torch.nn.functional as F

class LogisticRegression(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(LogisticRegression, self).__init__()
        self.linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        return F.sigmoid(self.linear(x))

训练逻辑回归模型

导入数据

MNIST 数据集包含 70,000 张手写数字图像，其中 60,000 张用于训练，10,000 张用于测试。我们可以使用 torchvision 库加载并预处理数据。

from torchvision import datasets, transforms

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())

定义训练参数

接下来，我们需要定义一些训练参数，例如学习率、批量大小和训练周期数。

learning_rate = 0.01
batch_size = 128
num_epochs = 10

构建数据加载器

数据加载器用于在训练和测试期间分批获取数据。

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

模型实例化和损失函数

我们实例化逻辑回归模型，并使用二元交叉熵损失函数。

model = LogisticRegression(784, 10)
loss_fn = nn.BCELoss()

训练循环

现在，我们可以执行训练循环，在每个周期中更新模型权重。

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

for epoch in range(num_epochs):
    for batch_idx, (data, target) in enumerate(train_loader):
        # 展平数据
        data = data.view(data.shape[0], -1)
        
        # 前向传播
        output = model(data)
        
        # 计算损失
        loss = loss_fn(output, target)
        
        # 反向传播
        loss.backward()
        
        # 优化器一步
        optimizer.step()
        
        # 清零梯度
        optimizer.zero_grad()

评估模型

训练后，我们可以通过计算准确率来评估模型的性能。

def accuracy(model, data_loader):
    correct = 0
    total = 0
    
    with torch.no_grad():
        for data, target in data_loader:
            # 展平数据
            data = data.view(data.shape[0], -1)
            
            # 前向传播
            output = model(data)
            
            # 获取预测标签
            pred = torch.argmax(output, dim=1)
            
            # 累加正确预测数量
            correct += (pred == target).sum().item()
            
            # 累加总样本数量
            total += target.shape[0]
    
    return correct / total

print(f"训练集准确率：{accuracy(model, train_loader) * 100:.2f}%")
print(f"测试集准确率：{accuracy(model, test_loader) * 100:.2f}%")