自定义数据集数字识别：PyTorch实现

本示例展示如何使用PyTorch构建一个神经网络模型，识别自定义的数字数据集，并包含训练和测试流程。

数据准备

首先，需要准备自定义的数字数据集，并确保其格式与MNIST数据集相同，即数据集应该包含'digital train'和'digital test'两个文件夹，每个文件夹下应该有数字命名的子文件夹，子文件夹中包含对应数字的图片。

代码实现

from torchvision.datasets import ImageFolder
from torchvision.transforms import Compose, ToTensor, Normalize
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam
import os
import torch
import numpy as np

BATCH_SIZE = 128
TEST_BATCH_SIZE = 1000

# 准备数据集
def get_dataloader(train, batch_size=BATCH_SIZE):
    transform_fn = Compose([
        ToTensor(),
        Normalize(mean=(0.1307,), std=(0.3081,))
    ])  # mean和std的形状与通道数相同

    dataset = ImageFolder(root='./digital train' if train else './digital test', transform=transform_fn)
    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    return data_loader


class MnistModel(nn.Module):
    def __init__(self):
        super(MnistModel, self).__init__()  # 继承
        self.fc1 = nn.Linear(1 * 28 * 28, 28)  # 参数是input和output的feature
        self.fc2 = nn.Linear(28, 10)

    def forward(self, input):
        # 1.进行形状的修改
        x = input.view([-1, 1 * 28 * 28])  # -1表示根据形状自动调整，也可以改为input.size(0)
        # 2.进行全连接的操作
        x = self.fc1(x)
        # 3.激活函数的处理
        x = F.relu(x)  # 形状没有变化
        # 4.输出层
        out = self.fc2(x)
        return F.log_softmax(out, dim=-1)


model = MnistModel()
optimizer = Adam(model.parameters(), lr=0.001)


def train(epoch):  # epoch表示几轮
    data_loader = get_dataloader(True)  # 获取数据加载器
    for idx, (input, target) in enumerate(data_loader):  # idx表示data_loader中的第几个数据，元组是data_loader的数据
        optimizer.zero_grad()  # 将梯度置0
        output = model(input)  # 调用模型，得到预测值
        loss = F.nll_loss(output, target)  # 调用损失函数，得到损失,是一个tensor
        loss.backward()  # 反向传播
        optimizer.step()  # 梯度的更新
        if idx % 10 == 0:
            print(epoch, idx, loss.item())
    # for是每一轮中的数据进行遍历


def test():
    loss_list = []
    acc_list = []
    test_dataloader = get_dataloader(train=False, batch_size=TEST_BATCH_SIZE)  # 获取测试集
    for idx, (input, target) in enumerate(test_dataloader):
        with torch.no_grad():  # 不计算梯度
            output = model(input)
            cur_loss = F.nll_loss(output, target)
            loss_list.append(cur_loss)
            # 计算准确率，output大小[batch_size,10] target[batch_size] batch_size是多少组数据，10列是每个数字概率
            pred = output.max(dim=-1)[-1]  # 获取最大值位置
            cur_acc = pred.eq(target).float().mean()
            acc_list.append(cur_acc)
    print('平均准确率：', np.mean(acc_list), '平均损失：', np.mean(loss_list))


test()
for i in range(1):  # 训练三轮
    train(i)
test()

使用方法

将上述代码保存为python文件，例如'digit_recognition.py'。
确保'digital train'和'digital test'文件夹在当前目录下。
运行'python digit_recognition.py'。

注意事项

确保数据集的格式与MNIST相同。
可以根据需要调整超参数，例如BATCH_SIZE、TEST_BATCH_SIZE、学习率等。
可以使用更复杂的模型结构来提升识别精度。

扩展

可以将代码改写为命令行工具，方便使用。
可以将模型保存并加载，方便以后使用。
可以将代码部署到服务器上，提供在线服务。