自己动手搭建卷积神经网络实现手写数字识别（Mnist数据集）

值得注意的是，在对数据进行预处理时需要使用transforms.RandomInvert()对图像的像素值进行翻转，因为Mnist数据集全都是黑底白字，直接使用其进行训练会导致模型无法识别白底黑字的情况（亲身经历）设计一个可以用于图像分类的卷积神经网络，首先明确输入的size大小为24*24的灰度图像，在此基础上使用公式。计算在前向传播的过程中，图片的尺寸变化，其中p是padding，k是kern

Logos515

559人浏览 · 2024-11-28 10:44:49

Logos515 · 2024-11-28 10:44:49 发布

一、配置环境

安装anaconda
创建anaconda虚拟环境
在环境下安装依赖（注意torch和cuda的版本要匹配，否则不能调用GPU）

二、模型创建

设计一个可以用于图像分类的卷积神经网络，首先明确输入的size大小为24*24的灰度图像，在此基础上使用公式

$n'=\lfloor\frac{(n+2p-k)}{s}\rfloor+1$

计算在前向传播的过程中，图片的尺寸变化，其中p是padding，k是kernel的边长，s是stride

我这里简单设计了一个卷积模型如下：

# 定义网络
class CustomNet(nn.Module):
    def __init__(self):
        super(CustomNet, self).__init__()
        self.features = nn.Sequential(
            # inputsize is 1 * 24 * 24
            nn.Conv2d(1, 16, 3, 2, 1),  # 输入通道 1（灰度图像），输出通道 16
            nn.ReLU(),
            # 16 * 12 * 12
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            # 32 * 12 * 12
            nn.MaxPool2d(2, 2),
            # 32 * 6 * 6
            nn.Conv2d(32, 20, 5, 1, 2),
            nn.ReLU()
            # 20 * 6 * 6
        )
        self.fc = nn.Sequential(
            nn.Linear(20 * 6 * 6, 30),
            nn.ReLU(),
            # 20 * 6 * 6 -> 30
            nn.Linear(30, 10),
            # 30 -> 10
            nn.Softmax(dim=1)  # 改为 dim=1，因为 Softmax 需要指定维度
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)  # 展平
        x = self.fc(x)
        return x

传播过程中图片的尺寸变化已经在注释里标出了，之所以要注意这一点是因为我们需要知道最后的全连接层的维度应该如何设计。

该模型的总参数量为50952，计算方法是：

对于卷积层，参数量为

$(c_{in}*h*w+1)*c_{out}$

其中h和w是卷积核的尺寸。

对于全连接层，参数量为

$(c_{in}+1)*c_{out}$

三、模型训练

我把自己的代码贴在下面：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 检查 GPU 可用性
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 调整 MNIST 数据集的图像大小为 24x24
transform = transforms.Compose([
    transforms.Resize((24, 24)),  # 调整到 24x24
    transforms.RandomInvert(),
    transforms.ToTensor()
])

# 加载 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)

# 定义网络
class CustomNet(nn.Module):
    def __init__(self):
        super(CustomNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, 3, 2, 1),  # 输入通道 1（灰度图像），输出通道 16
            nn.ReLU(),
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(32, 20, 5, 1, 2),
            nn.ReLU()
        )
        self.fc = nn.Sequential(
            nn.Linear(20 * 6 * 6, 30),
            nn.ReLU(),
            nn.Linear(30, 10),
            nn.Softmax(dim=1)  # 改为 dim=1，因为 Softmax 需要指定维度
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)  # 展平
        x = self.fc(x)
        return x

net = CustomNet().to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)

# 训练网络
num_epochs = 5

for epoch in range(num_epochs):
    net.train()
    running_loss = 0.0
    for images, labels in train_loader:
        # 将数据迁移到 GPU
        images, labels = images.to(device), labels.to(device)
        # 前向传播
        outputs = net(images)
        loss = criterion(outputs, labels)

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}')
    
# 测试网络
net.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:
        # 将数据迁移到 GPU
        images, labels = images.to(device), labels.to(device)

        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy: {100 * correct / total:.2f}%')

torch.save(net.state_dict(), 'model.pth')

值得注意的是，在对数据进行预处理时需要使用transforms.RandomInvert()对图像的像素值进行翻转，因为Mnist数据集全都是黑底白字，直接使用其进行训练会导致模型无法识别白底黑字的情况（亲身经历）

四、测试手写数字

除了测试集之外，可以自己手写一些数字让模型进行识别。

代码如下：

import torch
from torch import nn
from torchvision import transforms
from PIL import Image

# 定义网络
class CustomNet(nn.Module):
    def __init__(self):
        super(CustomNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, 3, 2, 1),  # 输入通道 1（灰度图像），输出通道 16
            nn.ReLU(),
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(32, 20, 5, 1, 2),
            nn.ReLU()
        )
        self.fc = nn.Sequential(
            nn.Linear(20 * 6 * 6, 30),
            nn.ReLU(),
            nn.Linear(30, 10),
            nn.Softmax(dim=1)  # 改为 dim=1，因为 Softmax 需要指定维度
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)  # 展平
        x = self.fc(x)
        return x
        
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
net = CustomNet().to(device)

net.load_state_dict(torch.load('model.pth', map_location=device))
net.eval()

# 定义图片预处理
transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),
    transforms.Resize((24,24)),
    transforms.ToTensor(),
])

# 加载本地图片
image_path = "7.png"  # 替换为您的图片路径
image = Image.open(image_path)
# 预处理图片
input_image = transform(image).unsqueeze(0).to(device)  # 增加批次维度

# 使用模型进行预测
with torch.no_grad():
    output = net(input_image)
    print(output)
    _, predicted = torch.max(output, 1)
    print(_, predicted)

# 输出预测结果
print(f"Predicted Label: {predicted.item()}")

测试之后发现其的确能够识别自己手写的数字。