基于pandas实现K折交叉验证数据集划分

k折交叉验证

将训练集分为k份，取其中一份成为验证集，剩下成为训练集，称为k折数据划分，使用得到的k个训练集、验证集，训练k个模型，k个模型做简单的模型融合，可提升模型稳定性、准确率。

数据介绍

数据共有四列，列名分别为id，title，content，id，抽样了20行方便进行试验。

代码实现

#以下代码得到上图的数据
train=pd.read_csv("data_0/train.csv",sep=",")
train=train[["id","title","content","label"]]
train=train[:20]
train.info()
#以上代码得到上图的数据

def k_fold_split(train_df,k):
    os.system("mkdir data")
    k_fold=[]
    index=set(range(train.shape[0]))
    for i in range(k):
        #防止所有数据不能整除k，最后将剩余的都放到最后一折
        if i==k-1:
            k_fold.append(list(index))
        else:
            tmp=random.sample(list(index),int(1.0/k*train.shape[0]))
            k_fold.append(tmp)
            index-=set(tmp)
    #将原始训练集划分为k个包含训练集和验证集的训练集，同时每个训练集中，训练集：验证集=k-1:1
    for i in range(k):
        print("第{}折........".format(i+1))
        tra=[]
        dev=k_fold[i]
        for j in range(k):
            if i!=j:
                tra+=k_fold[j]
        train.iloc[tra].to_csv("data/train_{}".format(i),sep=",",index=False)
        train.iloc[dev].to_csv("data/val_{}".format(i),sep=",",index=False)
    print("done!")


if __name__=="__main__":
    k_fold_split(train,5)

效果展示

打开文件验证是符合k折交叉验证数据划分原则。