pytorch实现手写数字图片识别

PyTorch是一个基于Python的科学计算库，它主要用于深度学习研究。在本文中，我们将介绍如何使用PyTorch实现手写数字图片识别。我们将分为两个部分，第一部分是数据预处理和模型训练，第二部分是模型测试和结果分析。

第一部分：数据预处理和模型训练

数据预处理

我们将使用MNIST数据集，该数据集包含60,000个训练图像和10,000个测试图像。每个图像都是28x28像素的灰度图像，表示手写数字0到9。我们将使用PyTorch内置的torchvision.datasets模块来加载MNIST数据集。下面是一个示例：

import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载MNIST数据集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# 定义数据加载器
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)

在这个示例中，我们首先定义了一个名为transform的数据预处理管道，该管道将图像转换为PyTorch张量，并对每个像素进行归一化。然后，我们使用datasets.MNIST类加载MNIST数据集，并使用transform参数将数据预处理应用于每个图像。我们还定义了两个数据加载器，一个用于训练数据，一个用于测试数据。每个数据加载器将数据分成大小为64的批次，并在每个时期之间打乱数据。

模型训练

我们将使用一个简单的卷积神经网络（CNN）来训练我们的模型。CNN是一种特殊的神经网络，它可以有效地处理图像数据。我们将使用PyTorch内置的torch.nn模块来定义CNN模型。下面是一个示例：

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

# 定义CNN模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.dropout2(x)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

# 定义优化器和损失函数
model = Net()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
criterion = nn.CrossEntropyLoss()

在这个示例中，我们首先定义了一个名为Net的CNN模型。该模型包含两个卷积层、两个dropout层和两个全连接层。我们使用nn.Conv2d类定义卷积层，使用nn.Dropout2d类定义dropout层，使用nn.Linear类定义全连接层。在forward方法中，我们使用F.relu函数对卷积层的输出进行非线性变换，使用F.max_pool2d函数对特征图进行下采样，使用torch.flatten函数将特征图展平为一维张量，使用F.log_softmax函数将输出转换为概率。然后，我们使用torch.optim.SGD类定义优化器，并使用nn.CrossEntropyLoss类定义损失函数。

接下来，我们将使用训练数据对模型进行训练。在每个时期中，我们将模型应用于每个批次的数据，并使用反向传播算法更新模型的权重。下面是一个示例：

# 训练模型
for epoch in range(10):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

    # 在测试数据上评估模型
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    print('Epoch: {} Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(
        epoch, test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

在这个示例中，我们首先使用model.train()方法将模型设置为训练模式。然后，我们使用enumerate函数遍历训练数据加载器中的每个批次，并使用optimizer.zero_grad()方法清除梯度。接下来，我们使用模型对数据进行预测，并使用criterion计算损失。然后，我们使用loss.backward()方法计算梯度，并使用optimizer.step()方法更新模型的权重。

在每个时期结束时，我们使用model.eval()方法将模型设置为评估模式，并使用torch.no_grad()上下文管理器禁用梯度计算。然后，我们使用测试数据加载器对模型进行评估，并计算平均损失和准确率。最后，我们打印出每个时期的平均损失和准确率。

第二部分：模型测试和结果分析

模型测试

在第一部分中，我们训练了一个CNN模型，并使用测试数据对其进行了评估。现在，我们将使用该模型对一些手写数字图像进行分类。下面是一个示例：

import matplotlib.pyplot as plt
import numpy as np

# 加载手写数字图像
image = plt.imread('test_image.png')
image = np.mean(image, axis=2)
image = 1 - image / 255.0
image = np.expand_dims(image, axis=0)
image = np.expand_dims(image, axis=0)

# 使用模型对图像进行分类
model.eval()
with torch.no_grad():
    output = model(torch.from_numpy(image).float())
    pred = output.argmax(dim=1, keepdim=True)

# 打印预测结果
print('Prediction:', pred.item())

在这个示例中，我们首先使用matplotlib.pyplot.imread函数加载手写数字图像，并使用numpy.mean函数将图像转换为灰度图像。然后，我们使用1 - image / 255.0将像素值归一化到0到1之间，并使用np.expand_dims函数将图像转换为PyTorch张量的形状。接下来，我们使用训练好的模型对图像进行分类，并使用output.argmax函数找到输出中最大值的索引。最后，我们打印出预测结果。

结果分析

我们可以使用混淆矩阵来分析模型的分类结果。混淆矩阵是一个表格，其中行表示实际类别，列表示预测类别。对角线上的元素表示正确分类的样本数，非对角线上的元素表示错误分类的样本数。下面是一个示例：

from sklearn.metrics import confusion_matrix

# 在测试数据上评估模型
model.eval()
test_loss = 0
correct = 0
y_true = []
y_pred = []
with torch.no_grad():
    for data, target in test_loader:
        output = model(data)
        test_loss += criterion(output, target).item()
        pred = output.argmax(dim=1, keepdim=True)
        correct += pred.eq(target.view_as(pred)).sum().item()
        y_true.extend(target.numpy())
        y_pred.extend(pred.numpy().flatten())

test_loss /= len(test_loader.dataset)
print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(
    test_loss, correct, len(test_loader.dataset),
    100. * correct / len(test_loader.dataset)))

# 计算混淆矩阵
cm = confusion_matrix(y_true, y_pred)
print(cm)

在这个示例中，我们首先定义了两个空列表y_true和y_pred，用于存储测试数据的实际标签和预测标签。然后，我们使用测试数据加载器对模型进行评估，并计算平均损失和准确率。在评估过程中，我们使用extend方法将每个批次的实际标签和预测标签添加到列表中。最后，我们使用sklearn.metrics.confusion_matrix函数计算混淆矩阵，并打印出结果。

混淆矩阵可以帮助我们了解模型在每个类别上的性能。例如，如果模型在某个类别上的准确率很低，我们可以尝试增加该类别的训练数据或调整模型的超参数。

本站文章如无特殊说明，均为本站原创，如若转载，请注明出处：pytorch实现手写数字图片识别 - Python技术站