社区首页 >专栏 >基于OCR模型的训练数据划分教程

基于OCR模型的训练数据划分教程

原创

IT蜗壳-Tango

发布于 2024-08-04 14:46:44

26403

代码可运行

文章被收录于专栏：机器学习机器学习

运行总次数：3

代码可运行

在训练OCR（光学字符识别）模型时，数据集的划分是至关重要的步骤。合理的划分能确保模型的泛化能力，即在未见过的数据上仍能表现良好。本文将详细介绍如何划分训练集、验证集和测试集，确保模型的性能和可靠性。

1. 数据集准备

在开始数据集划分之前，首先需要准备好原始数据集。OCR任务的数据集通常由带有文字的图像及其对应的标签（文本）组成。一个典型的数据集可能包含成千上万张图像，涵盖各种字体、语言和文本布局。

1.1 数据收集

多样性：确保数据集涵盖不同的字体、大小、语言、背景和噪声情况。
标注质量：每张图像都应有精确的文本标签，错误或不完整的标签会影响模型的训练效果。

2. 数据集划分

数据集通常划分为三个部分：训练集（Training Set）、验证集（Validation Set）和测试集（Test Set）。

2.1 训练集

训练集用于训练模型，是数据集中最大的一部分。一般来说，训练集占整个数据集的60%到80%。训练集中的样本应尽可能全面，涵盖所有可能的场景和变体，以便模型能够学习到足够的信息。

2.2 验证集

验证集用于调优模型超参数以及选择最佳模型。通常占数据集的10%到20%。验证集应与训练集保持一致性，但又不能完全相同，以避免过拟合。通过在验证集上的表现，我们可以调整模型的结构和参数，确保模型的泛化能力。

2.3 测试集

测试集用于评估最终模型的性能，通常占数据集的10%到20%。测试集应在训练过程中完全隔离，不能用于任何模型调整。只有在训练和验证完成后，才能使用测试集进行评估，以提供一个真实的性能衡量标准。

3. 数据集划分策略

3.1 随机划分

最简单的方法是随机划分数据集。假设有10000张图像，可以随机抽取6000-8000张作为训练集，1000-2000张作为验证集，1000-2000张作为测试集。

from sklearn.model_selection import train_test_split

# 假设 images 是图像列表，labels 是对应的标签列表
train_images, test_images, train_labels, test_labels = train_test_split(images, labels, test_size=0.2, random_state=42)
train_images, val_images, train_labels, val_labels = train_test_split(train_images, train_labels, test_size=0.25, random_state=42)
# 最终划分比例为：训练集 60%，验证集 20%，测试集 20%

3.2 分层抽样

对于不平衡数据集，分层抽样可以确保每个类别在训练集、验证集和测试集中都有相同比例的样本。这对于OCR模型特别重要，因为不同字符、字体和语言的分布可能非常不均匀。

from sklearn.model_selection import StratifiedShuffleSplit

# 分层抽样划分
sss = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
for train_index, test_index in sss.split(images, labels):
    train_images, test_images = images[train_index], images[test_index]
    train_labels, test_labels = labels[train_index], labels[test_index]

sss = StratifiedShuffleSplit(n_splits=1, test_size=0.25, random_state=42)
for train_index, val_index in sss.split(train_images, train_labels):
    train_images, val_images = train_images[train_index], train_images[val_index]
    train_labels, val_labels = train_labels[train_index], train_labels[val_index]
# 最终划分比例为：训练集 60%，验证集 20%，测试集 20%

3.3 时间序列划分

如果数据集具有时间相关性（例如OCR任务中的连续扫描页），应根据时间顺序进行划分，确保训练集、验证集和测试集都涵盖不同时期的数据，避免模型只在特定时间段的数据上表现良好。

4. 数据增强

在数据集划分后，可以对训练集进行数据增强，以增加数据的多样性。常用的增强方法包括旋转、缩放、翻转、添加噪声等。这些操作可以帮助模型更好地泛化，减少过拟合。

from tensorflow.keras.preprocessing.image import ImageDataGenerator

datagen = ImageDataGenerator(
    rotation_range=10,
    width_shift_range=0.1,
    height_shift_range=0.1,
    shear_range=0.1,
    zoom_range=0.1,
    horizontal_flip=True,
    fill_mode='nearest')

# 训练集增强示例
datagen.fit(train_images)

5. 实践案例

假设我们有一个包含10000张图像的OCR数据集，标签包括英文、数字和一些特殊字符。我们可以使用上述方法将数据集划分为：

训练集：6000张
验证集：2000张
测试集：2000张

通过分层抽样确保每个字符类别在三个子集中都有相同比例的样本。然后对训练集进行数据增强，增加数据的多样性。

# 数据集划分
from sklearn.model_selection import train_test_split

# 随机划分数据集
train_images, test_images, train_labels, test_labels = train_test_split(images, labels, test_size=0.2, random_state=42)
train_images, val_images, train_labels, val_labels = train_test_split(train_images, train_labels, test_size=0.25, random_state=42)

# 数据增强
from tensorflow.keras.preprocessing.image import ImageDataGenerator

datagen = ImageDataGenerator(
    rotation_range=10,
    width_shift_range=0.1,
    height_shift_range=0.1,
    shear_range=0.1,
    zoom_range=0.1,
    horizontal_flip=True,
    fill_mode='nearest')

datagen.fit(train_images)

# 训练模型示例
model.fit(datagen.flow(train_images, train_labels, batch_size=32), epochs=50, validation_data=(val_images, val_labels))