【模型+代码/保姆级教程】使用Pytorch实现手写汉字识别

该教程详细介绍了如何使用Pytorch框架和EfficientNetV2模型进行手写汉字识别,包括数据预处理、模型训练、评估和推理的全过程。作者提供了预训练模型和预处理数据集的下载链接,以及项目源码,适合新手入门。
摘要由CSDN通过智能技术生成

前言


参考文章:

最初参考的两篇:
【Pytorch】基于CNN手写汉字的识别
「Pytorch」CNN实现手写汉字识别(数据集制作,网络搭建,训练验证测试全部代码)
模型:
EfficientNetV2网络详解
数据集(不必从这里下载,可以看一下它的介绍):
CASIA Online and Offline Chinese Handwriting Databases

鉴于已经3202年了,大模型都出来了,网上还是缺乏汉字识别这种基础神经网络的能让新手直接上手跑通的手把手教程,我就斗胆自己写一篇好了。

本文的主要特点:

  1. 使用EfficientNetV2模型真正实现3755类汉字识别

  2. Demo开源

  3. 预训练模型可以下载

  4. 预制数据集,无需处理直接使用

数据集


使用中科院制作的手写汉字数据集,链接直达官网,所以我这里不多介绍,只有满腔敬意。

上面参考的博客可能要你自己下载之后按照它的办法再预处理一下,但是在这个环节出现问题的朋友挺多,我把预处理的数据已经传到【鼠盘】,有人反映说这个有损坏了,但是我自己用bandizip智能解压一切正常,没损坏。如果还是报错,则使用这个别人传的【百度网盘】

预训练模型已经上传了(后面有链接),但是如果想自己训一下,就需要下载这个数据集,解压到项目结构里的data文件夹如下所示

data文件夹和log文件夹需要自己建。

项目结构


完整源代码:【项目源码

目录结构

重点注意data文件夹的结构,不要把数据集放错位置了或者多嵌套了文件夹

├─Chinese_Character_Rec(项目)
│ ├─asserts
│ │ ├─*.png
│ ├─char_dict
│ ├─Data.py
│ ├─EfficientNetV2
│ │ ├─demo.py
│ │ ├─EffNetV2.py
│ │ ├─Evaluate.py
│ │ ├─model.py
│ │ └─Train.py
│ ├─Utils.py
│ ├─VGG19
│ │ ├─demo.py
│ │ ├─Evaluate.py
│ │ ├─model.py
│ │ ├─Train.py
│ │ └─VGG19.py
│ └─README.md
├─data(数据集)
│ ├─test(测试数据集)
│ │ ├─00000
│ │ ├─00001
│ │ └─...
│ ├─test.txt(程序生成)
│ ├─train(训练数据集)
│ │ ├─00000
│ │ ├─00001
│ │ └─ ...
│ └─train.txt(程序生成)
├─log(模型参数存放位置)
    ├─log1.pth
    └─…

神经网络模型


预训练模型【参数链接】(包含vgg19和efficientnetv2)

请将.pth文件重命名为log+数字.pth的格式,例如log1.pth,放入log文件夹。方便识别和retrain。

VGG19

这里先后用了两种神经网络,我先用VGG19试了一下,分类前1000种汉字。训得有点慢,主要还是这模型有点老了,参数量也不小。而且要改到3755类的话还用原参数的话就很难收敛,也不知道该怎么调参数了,估计调好了也会规模很大,所以这里VGG19模型的版本只能分类1000种,就是数据集的前1000种(准确率>92%)。

EfficientNetV2

这个模型很不错,主要是卷积层的部分非常有效,参数量也很少。直接用small版本去分类3755个汉字,半小时就收敛得差不多了。所以本文用来实现3755类汉字的模型就是EfficientNetV2(准确率>89%),后面的教程都是基于这个,VGG19就不管了,在源码里感兴趣的自己看吧。

以下代码不用自己写,前面已经给出完整源代码了,下面的教程是结合源码的讲解而已。

运行环境


显存>=4G(与batchSize有关,batchSize=512时显存占用4.8G;如果是256或者128,应该会低于4G,虽然会导致训得慢一点)

内存>=16G(训练时不太占内存,但是刚开始加载的时候会突然占一下,如果小于16G还是怕爆)

如果你没有安装过Pytorch,啊,我也不知道怎么办,你要不就看看安装Pytorch的教程吧。(总体步骤是,有一个不太老的N卡,先去驱动里看看cuda版本,安装合适的CUDA,然后根据CUDA版本去pytorch.org找到合适的安装指令,然后在本地pip install)

以下是项目运行环境,我是3060 6G,CUDA版本11.6

这个约等号不用在意,可以都安装最新版本,反正我这里应该没用什么特殊的API


torch~=1.12.1+cu116
torchvision~=0.13.1+cu116
Pillow~=9.3.0

数据集准备


首先定义classes_txt方法在Utils.py中(不是我写的,是CSDN那两篇博客的,MyDataset同):

生成每张图片的路径,存储到train.txt或test.txt。方便训练或评估时读取数据


def classes_txt(root, out_path, num_class=None):
    dirs = os.listdir(root)
    if not num_class:
        num_class = len(dirs)

    with open(out_path, 'w') as f:
        end = 0
        if end < num_class - 1:
            dirs.sort()
            dirs = dirs[end:num_class]
            for dir1 in dirs:
                files = os.listdir(os.path.join(root, dir1))
                for file in files:
                    f.write(os.path.join(root, dir1, file) + '\n')

定义Dataset类,用于制作数据集,为每个图片加上对应的标签,即图片所在文件夹的代号


class MyDataset(Dataset):
    def __init__(self, txt_path, num_class, transforms=None):
        super(MyDataset, self).__init__()
        images = []
        labels = []
        with open(txt_path, 'r') as f:
            for line in f:
                if int(line.split('\\')[1]) >= num_class: # 超出规定的类,就不添加,例如VGG19只添加了1000类
                    break
                line = line.strip('\n')
                images.append(line)
                labels.append(int(line.split('\\')[1]))
        self.images = images
        self.labels = labels
        self.transforms = transforms

    def __getitem__(self, index):
        image = Image.open(self.images[index]).convert('RGB')
        label = self.labels[index]
        if self.transforms is not None:
            image = self.transforms(image)
        return image, label

    def __len__(self):
        return len(self.labels)

入口


我把各种超参都放在了args里方便改,请根据实际情况自行调整。这套defaults就是我训练这个模型时使用的超参,图片size默认32是因为我显存太小辣!!但是数据集给的图片大小普遍不超过64,如果想训得更精确,可以试试64*64的大小。

如果你训练时爆mem,请调小batch_size,试试256,128,64,32


parser = argparse.ArgumentParser(description='EfficientNetV2 arguments')
parser.add_argument('--mode', dest='mode', type=str, default='demo', help='Mode of net')
parser.add_argument('--epoch', dest='epoch', type=int, default=50, help='Epoch number of training')
parser.add_argument('--batch_size', dest='batch_size', type=int, default=512, help='Value of batch size')
parser.add_argument('--lr', dest='lr', type=float, default=0.0001, help='Value of lr')
parser.add_argument('--img_size', dest='img_size', type=int, default=32, help='reSize of input image')
parser.add_argument('--data_root', dest='data_root', type=str, default='../../data/', help='Path to data')
parser.add_argument('--log_root', dest='log_root', type=str, default='../../log/', help='Path to model.pth')
parser.add_argument('--num_classes', dest='num_classes', type=int, default=3755, help='Classes of character')
parser.add_argument('--demo_img', dest='demo_img', type=str, default='../asserts/fo2.png', help='Path to demo image')
args = parser.parse_args()


if __name__ == '__main__':
    if not os.path.exists(args.data_root + 'train.txt'): # 只生成一次
        classes_txt(args.data_root + 'train', args.data_root + 'train.txt', args.num_classes)
    if not os.path.exists(args.data_root + 'test.txt'): # 只生成一次
        classes_txt(args.data_root + 'test', args.data_root + 'test.txt', args.num_classes)

    if args.mode == 'train':
        train(args)
    elif args.mode == 'evaluate':
        evaluate(args)
    elif args.mode == 'demo':
        demo(args)
    else:
        print('Unknown mode')

训练


在前面CSDN博客的基础上,增加了lr_scheduler自行调整学习率(如果连续2个epoch无改进,就调小lr到一半),增加了连续训练的功能:

先在log文件夹下寻找是否存在参数文件,如果没有,就认为是初次训练;如果有,就找到后缀数字最大的log.pth,在这个基础上继续训练,并且每训练完一个epoch,就保存最新的log.pth,代号是上一次的+1。这样可以多次训练,防止训练过程中出错,参数文件损坏前功尽弃。

其中has_log_file和find_max_log在Utils.py中有定义。


def train(args):
    print("===Train EffNetV2===")
    # 归一化处理,不一定要这样做,看自己的需求,只是预训练模型的训练是这样设置的
    transform = transforms.Compose(
        [transforms.Resize((args.img_size, args.img_size)), transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
         transforms.ColorJitter()])  

    train_set = MyDataset(args.data_root + 'train.txt', num_class=args.num_classes, transforms=transform)
    train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True)
    device = torch.device('cuda:0')
    # 加载模型
    model = efficientnetv2_s(num_classes=args.num_classes)
    model.to(device)
    model.train()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=args.lr)
    # 学习率调整函数,不一定要这样做,可以自定义
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=2, factor=0.5)
    print("load model...")
    
    # 加载最近保存了的参数
    if has_log_file(args.log_root):
        max_log = find_max_log(args.log_root)
        print("continue training with " + max_log + "...")
        checkpoint = torch.load(max_log)
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        loss = checkpoint['loss']
        epoch = checkpoint['epoch'] + 1
    else:
        print("train for the first time...")
        loss = 0.0
        epoch = 0

    while epoch < args.epoch:
        running_loss = 0.0
        for i, data in enumerate(train_loader):
            inputs, labels = data[0].to(device), data[1].to(device)
            optimizer.zero_grad()
            outs = model(inputs)
            loss = criterion(outs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
            if i % 200 == 199:
                print('epoch %5d: batch: %5d, loss: %8f, lr: %f' % (
                    epoch + 1, i + 1, running_loss / 200, optimizer.state_dict()['param_groups'][0]['lr']))
                running_loss = 0.0

        scheduler.step(loss)
        # 每个epoch结束后就保存最新的参数
        print('Save checkpoint...')
        torch.save({'epoch': epoch,
                    'model_state_dict': model.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'loss': loss},
                   args.log_root + 'log' + str(epoch) + '.pth')
        print('Saved')
        epoch += 1

    print('Finish training')

评估


跑测试集,算总体准确率。有一点不完善,就是看不到每一个类具体的准确率。我的预训练模型其实感觉有几类是过拟合的,但是我懒得调整了。


def evaluate(args):
    print("===Evaluate EffNetV2===")
    # 这个地方要和train一致,不过colorJitter可有可无
    transform = transforms.Compose(
        [transforms.Resize((args.img_size, args.img_size)), transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
         transforms.ColorJitter()])

    model = efficientnetv2_s(num_classes=args.num_classes)
    model.eval()
    if has_log_file(args.log_root):
        file = find_max_log(args.log_root)
        print("Using log file: ", file)
        checkpoint = torch.load(file)
        model.load_state_dict(checkpoint['model_state_dict'])
    else:
        print("Warning: No log file")

    model.to(torch.device('cuda:0'))
    test_loader = DataLoader(MyDataset(args.data_root + 'test.txt', num_class=args.num_classes, transforms=transform),batch_size=args.batch_size, shuffle=False)
    total = 0.0
    correct = 0.0
    print("Evaluating...")
    with torch.no_grad():
        for i, data in enumerate(test_loader):
            inputs, labels = data[0].cuda(), data[1].cuda()
            outputs = model(inputs)
            _, predict = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predict == labels).sum().item()
    acc = correct / total * 100
    print('Accuracy'': ', acc, '%')

推理


输入文字图片,输出识别结果:

其中char_dict就是每个汉字在数据集里的代号对应的gb2312编码,这个模型的输出结果是它在数据集里的代号,所以要查这个char_dict来获取它对应的汉字。


def demo(args):
    print('==Demo EfficientNetV2===')
    print('Input Image: ', args.demo_img)
    # 这个地方要和train一致,不过colorJitter可有可无
    transform = transforms.Compose(
        [transforms.Resize((args.img_size, args.img_size)), transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
    img = Image.open(args.demo_img)
    img = transform(img)
    img = img.unsqueeze(0) # 增维
    model = efficientnetv2_s(num_classes=args.num_classes)
    model.eval()
    if has_log_file(args.log_root):
        file = find_max_log(args.log_root)
        print("Using log file: ", file)
        checkpoint = torch.load(file)
        model.load_state_dict(checkpoint['model_state_dict'])
    else:
        print("Warning: No log file")

    with torch.no_grad():
        output = model(img)
    _, pred = torch.max(output.data, 1)
    f = open('../char_dict', 'rb')
    dic = pickle.load(f)
    for cha in dic:
        if dic[cha] == int(pred):
            print('predict: ', cha)
    f.close()

例如输入图片为:

程序运行结果:

其他说明


如遇到Module not found之类的错,重新写一下import,从实际的位置导入。

这个模型我通过ChaquoPy尝试移植到了Android平台,不过效果一般,我也没好好做:手写汉字识别APP,借用开源手写板

另外,这个模型对于太细太黑的字体,准确度貌似不是很好,可能还是有点过拟合了。建议输入的图片与数据集的风格靠拢,黑色尽量浅一点,线不要太细。

B站同步文章:(【模型+代码/保姆级教程】使用Pytorch实现手写汉字识别 - 哔哩哔哩

2023年9月更新:本项目已不再做,只是本人的学习实践,教程只是帮你跑通一个简单有效果的深度学习,事实上移动端和PC端应该都有比EffNetV2更加合适的模型,且需要仔细设定学习策略。有链接失效可以B站私信扣我,我过来补链接。

手写数字识别是一项常见的机器学习任务,可以使用PyTorch库来实现。首先,我们需要准备手写数字的数据集,比如MNIST数据集,它包含大量的手写数字图片和对应的标签。然后,我们可以使用PyTorch提供的数据加载和预处理工具来加载数据集,并将图片转换成张量,标签转换成对应的数字。 接下来,我们可以构建一个卷积神经网络(CNN)模型来进行手写数字识别。在PyTorch中,我们可以使用nn.Module类来定义神经网络模型,通过继承这个类并重写forward方法来定义模型的前向传播过程。我们可以使用一些常见的卷积层、池化层和全连接层来构建我们的CNN模型,同时也可以添加一些激活函数和正则化方法来提高模型的性能。 在模型定义好之后,我们可以使用PyTorch提供的优化器和损失函数来进行模型训练。我们可以使用反向传播算法来更新模型的参数,通过调整参数来最小化损失函数。在训练过程中,我们也可以使用一些常见的训练技巧,比如学习率衰减、批量归一化等来提高模型的收敛速度和准确率。 最后,我们可以使用训练好的模型来进行手写数字识别。我们可以将手写数字图片输入到模型中,模型会输出对应的数字标签,我们可以对这个标签进行解码,得到我们识别出的手写数字。通过不断优化模型和参数,我们可以实现准确率较高的手写数字识别模型
评论 29
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

Katock_Cricket

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值