9 数据集准备之数据划分

在微调大模型（LLM）之前，正确的数据集准备是至关重要的一步。在上一篇中，我们探讨了数据格式化的过程，包括如何将原始数据转换为模型可接受的格式。本文将聚焦于数据划分，讨论如何将数据集有效分配为训练集、验证集和测试集，以确保模型能够在实际应用中表现良好。

数据划分的重要性

数据划分的主要目的是为了确保模型在不同阶段的训练和评估中都有充足的数据支持。一般来说，我们将数据集划分为以下几个部分：

1. 训练集（Training Set）：用于模型的训练。通常占总数据集的70-80%。
2. 验证集（Validation Set）：在训练过程中用于调优模型的参数。通常占总数据集的10-15%。
3. 测试集（Test Set）：用于模型训练完成后的评估，确保模型在未见数据上的性能。通常占总数据集的10-15%。

数据划分策略

在进行数据划分时，有多种策略可以选择，下面我们将讨论两种常用的方法。

1. 随机划分

最简单、最常用的方法是随机划分数据集。可以使用Python的scikit-learn库中的train_test_split函数来实现。以下是一个示例代码：

import pandas as pd
from sklearn.model_selection import train_test_split

# 假设我们有一个名为data的数据框，包含了原始数据
data = pd.read_csv('your_dataset.csv')

# 将数据随机划分为训练集和测试集（80% 训练集，20% 测试集）
train_data, test_data = train_test_split(data, test_size=0.2, random_state=42)

# 接着再从训练集中划分出验证集（20% 的训练集，80% 的训练集）
train_data, val_data = train_test_split(train_data, test_size=0.2, random_state=42)

print(f"训练集大小: {train_data.shape}")
print(f"验证集大小: {val_data.shape}")
print(f"测试集大小: {test_data.shape}")

在这个例子中，原始数据集被分为80%的训练集，20%的测试集；然后训练集又被划分为80%用作训练，20%用作验证。

2. 分层划分

如果数据集中的类别分布不均，例如在分类任务中，一些类别的样本数量远多于其他类别，使用简单的随机划分可能会导致训练、验证、测试集中类别不均衡。

此时我们可以选择分层划分（Stratified Split），通过指定stratify参数在train_test_split中实现：

# 假设'labels'是我们的目标变量
train_data, test_data = train_test_split(data, test_size=0.2, stratify=data['labels'], random_state=42)
train_data, val_data = train_test_split(train_data, test_size=0.2, stratify=train_data['labels'], random_state=42)

这样确保了每个数据集中的类别比例大致与原始数据集一致。

划分后的数据集检查

完成数据划分后，我们需要检查每个数据集的大小和类别分布，以确保划分的合理性。以下代码展示了如何打印各个数据集合的类别分布：

print("训练集类别分布：")
print(train_data['labels'].value_counts(normalize=True))

print("验证集类别分布：")
print(val_data['labels'].value_counts(normalize=True))

print("测试集类别分布：")
print(test_data['labels'].value_counts(normalize=True))

总结

本章中我们探讨了数据划分的重要性和常用策略，包括随机划分和分层划分。数据划分不仅影响模型的训练效果，更直接关系到模型在实际应用中的表现。因此，采用合理的划分方法，确保训练、验证和测试集的代表性是至关重要的。

在下一篇中，我们将讨论如何选择合适的预训练模型及其基础知识，以便于在微调阶段充分利用这些模型的潜力。请继续关注我们的系列教程，了解更多关于大模型微调的知识。