机器学习预测-CNN手写字识别

news2025/3/15 17:34:03

介绍

这段代码是使用PyTorch实现的卷积神经网络（CNN），用于在MNIST数据集上进行图像分类。让我一步步解释：

导入库：代码导入了必要的库，包括PyTorch（torch）、神经网络模块（torch.nn）、函数模块（torch.nn.functional）、图像数据集（torchvision）以及数据处理（torch.utils.data）和可视化（matplotlib.pyplot）的工具。
设置超参数：定义了超参数，如批大小（Batch_size）、epoch数量（Epoch）和学习率（Lr）。
加载MNIST数据集：使用torchvision.datasets.MNIST加载MNIST数据集。该数据集包含了0到9的手写数字的灰度图像。transform=torchvision.transforms.ToTensor()将PIL图像转换为PyTorch张量。
可视化样本数据：打印数据集的大小，并显示数据集中的第一张图像及其相应的标签。
准备测试数据：准备测试数据与训练数据类似。加载MNIST测试数据集，并选择前2000个图像进行测试。
创建数据加载器：使用torch.utils.data.DataLoader创建训练数据的数据加载器。它有助于在训练过程中对数据进行分批和混洗。
定义CNN架构：通过子类化nn.Module来定义CNN类。该架构包括两个卷积层（self.con1、self.con2），后面跟有ReLU激活函数和最大池化层。卷积层的输出被展平并馈入全连接层（self.out），产生最终输出。
初始化CNN：创建CNN类的实例。
定义损失函数和优化器：使用交叉熵损失（nn.CrossEntropyLoss）作为损失函数，使用随机梯度下降（torch.optim.SGD）作为优化器。
训练CNN：在指定的epoch数量循环内训练模型。在循环内，将训练数据通过模型，计算损失，进行梯度反向传播，并由优化器更新模型参数。
测试模型：每50次迭代训练时，对测试数据集进行评估。将测试预测与真实标签进行比较，计算准确率。
打印结果：训练结束后，打印模型预测及前10个测试样本的真实标签。

总的来说，这段代码训练了一个CNN模型，用于在MNIST数据集上对手写数字进行分类，并在单独的测试数据集上评估其性能。

代码

# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torch.utils.data as Data
import matplotlib.pyplot as plt

# define hyper parameters
Batch_size = 100
Epoch = 1
Lr = 0.5
#DOWNLOAD_MNIST = True # 若没有数据，用此生成数据

# define train data and test data
train_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=True,
    download=False,
    transform=torchvision.transforms.ToTensor()
)
print(train_data.data.size())
print(train_data.targets.size())
print(train_data.data[0])
# 画一个图片显示出来
plt.imshow(train_data.data[0].numpy(),cmap='gray')
plt.title('%i'%train_data.targets[0])
plt.show()
# print(train_data.data.shape)           # torch.Size([60000, 28, 28])
# print(train_data.targets.size())        # torch.Size([60000])
# print(train_data.data[0].size())        # torch.Size([28, 28])
# plt.imshow(train_data.data[0].numpy(), cmap='gray')
# plt.show()
test_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=False,
    # transform=torchvision.transforms.ToTensor()
)
test_x = torch.unsqueeze(test_data.data, dim=1).type(torch.FloatTensor)[:2000]
test_y = test_data.targets[:2000]
# print(test_x.shape)                         # torch.Size([2000, 1, 28, 28])
# print(test_y.shape)                         # torch.Size([2000])
train_loader = Data.DataLoader(
    dataset=train_data,
    shuffle=True,
    batch_size=Batch_size,
)

# define network structure
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.con1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.con2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.out = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x):
        x = self.con1(x)            # (batch, 16, 14, 14)
        x = self.con2(x)            # (batch, 32, 7, 7)
        x = x.view(x.size(0), -1)
        out = self.out(x)             # (batch_size, 10)
        return out

cnn = CNN()
# print(cnn)
optimizer = torch.optim.SGD(cnn.parameters(), lr=Lr)
loss_fun = nn.CrossEntropyLoss()

for epoch in range(Epoch):
    for i, (x, y) in enumerate(train_loader):
        output = cnn(x)
        loss = loss_fun(output, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if i % 50 == 0:
            test_output = torch.max(cnn(test_x), dim=1)[1]
            loss = loss_fun(cnn(test_x), test_y).item()
            accuracy = torch.sum(torch.eq(test_output, test_y)).item() / test_y.numpy().size
            print('Epoch:', Epoch, '|loss:%.4f' % loss, '|accuracy:%.4f' % accuracy)

print('real value', test_data.targets[: 10].numpy())
print('train value', torch.max(cnn(test_x)[: 10], dim=1)[1].numpy())