python 划分训练集——K折交叉验证

最新推荐文章于 2025-09-11 16:52:12 发布

原创最新推荐文章于 2025-09-11 16:52:12 发布 · 4.2k 阅读

15 ·

CC 4.0 BY-SA版权

我们不生产知识，我们只是互联网的搬运工

编程语言同时被 2 个专栏收录

240 篇文章

订阅专栏

机器学习

136 篇文章

订阅专栏

本文介绍如何使用Python和sklearn库中的StratifiedKFold进行五折交叉验证的数据集划分，确保训练集和测试集的类别分布一致，适用于监督学习任务。

首先来生成一个训练集

import pandas as pd
import numpy as np
from sklearn.model_selection import StratifiedKFold

filename_label = {'filename':[str(i)+'.jpg' for i in range(100)], 'label':[np.random.randint(0,5) for i in range(100)]}
train = pd.DataFrame(filename_label)

print(train['label'].value_counts())
'''
2    23
1    23
0    20
4    18
3    16
Name: label, dtype: int64
'''
train.head(10)

在这里插入图片描述

接下来使用 sklearn.model_selection.StratifiedKFold，把这个 CSV 文件分成 2K 个文件，即 K 个训练集加 K 个测试集：

n_splits = 5  # K
x = train['filename'].values
y = train['label'].values
skf = StratifiedKFold(n_splits=n_splits, random_state=42, shuffle=True)

for index,(train_index,test_index) in enumerate(skf.split(x,y), start=1):
    res_train = pd.DataFrame()
    res_train['filename'] = train['filename'].iloc[train_index]
    res_train['label'] = train['label'].iloc[train_index]
    res_train.to_csv("train_{}.csv".format(index),index=False)

    res_train = pd.DataFrame()
    res_train['filename'] = train['filename'].iloc[test_index]
    res_train['label'] = train['label'].iloc[test_index]
    res_train.to_csv("test_{}.csv".format(index),index=False)