数据集进行测试集划分时按照标签的比例划分
时间: 2024-09-24 11:09:03 浏览: 32
用pandas划分数据集实现训练集和测试集
5星 · 资源好评率100%
当你需要在数据集中划分测试集时,保持标签比例通常是统计学上的一种良好实践,特别是对于不平衡数据集,比如某些类别的样本数量远多于其他类别。这种做法可以保证你在测试阶段能够得到类似真实场景下的性能评估。
例如,如果你想将数据集划分为70%的训练集和30%的测试集,并且希望训练集和测试集的各类别分布尽可能一致,你可以采用以下步骤:
1. **确定类别和样本总数**:首先了解数据集中所有类别的名称及其样本数量。
2. **计算各类别在总数据集中的占比**:这有助于了解哪些类别占比较大,哪些较小。
3. **随机抽样**:针对每一种类别,从该类别的样本中抽取相应比例的部分作为测试集的一部分。例如,如果一个类别的样本在总样本中占比为15%,那么你需要从这个类别中随机选择15%的样本放入测试集。
4. **合并**:将所有类别抽样的样本合并起来,就得到了按标签比例划分好的测试集。
5. **剩余数据为训练集**:剩下的样本即为训练集。
Python中有许多库可以帮助你完成这样的操作,例如`train_test_split`函数可以从`sklearn.model_selection`模块中找到,它支持按标签比例划分。
```python
from sklearn.model_selection import train_test_split
# 假设df是包含特征和标签的数据框,标签列名为'label'
label_counts = df['label'].value_counts() # 获取每个类别的计数
total_samples = df.shape[0]
test_size_per_label = test_size / label_counts.sum() # 根据总体大小和期望的测试集比例计算每个类别的测试样本数
test_data = {}
for label, count in label_counts.items():
test_samples = int(count * test_size_per_label)
test_indices = np.random.choice(df[df['label'] == label].index, test_samples, replace=False)
test_data[label] = df.loc[test_indices]
# 合并所有类别的测试数据
test_df = pd.concat(test_data.values(), ignore_index=True)
# 训练集 = 原始数据 - 测试数据
train_df = df.drop(test_df.index)
# 划分训练集和测试集
train_X, test_X, train_y, test_y = train_test_split(train_df.drop('label', axis=1), train_df['label'], test_size=test_size, stratify=train_df['label'])
```
阅读全文