首页 > 解决方案 > 如何在python中将数据表数据框拆分为训练和测试数据集

问题描述

我正在使用数据表数据。如何将数据框拆分为训练和测试数据集?
与 pandas 数据框类似,我尝试train_test_split(dt_df,classes)从 sklearn.model_selection 使用,但它不起作用并且出现错误。

import datatable as dt
import numpy as np
from sklearn.model_selection import train_test_split

dt_df = dt.fread(csv_file_path)
classe = dt_df[:, "classe"])
del dt_df[:, "classe"])

X_train, X_test, y_train, y_test = train_test_split(dt_df, classe, test_size=test_size)

我收到以下错误:TypeError: Column selector must be an integer or a string, not <class 'numpy.ndarray'>

我通过将数据框转换为 numpy 数组来尝试解决方法:

classe = np.ravel(dt_df[:, "classe"])
dt_df = dt_df.to_numpy()

就像它一样工作,但是,我不知道是否有一种方法可以train_test_split像熊猫数据框一样正常工作。

编辑 1: csv 文件包含作为列字符串,并且值是无符号整数。使用print(dt_df)我们得到:

     | CCC CCG CCU CCA CGC CGG CGU CGA CUC CUG …  
---- + --- --- --- --- --- --- --- --- --- ---     
   0 | 0 0 0 0 2 0 1 0 0 1 …  
   1 | 0 0 0 0 1 0 2 1 0 1 …  
   2 | 0 0 0 1 1 0 1 0 1 2 …  
   3 | 0 0 0 1 1 0 1 0 1 2 …  
   4 | 0 0 0 1 1 0 1 0 1 2 …  
   5 | 0 0 0 1 1 0 1 0 1 2 …  
   6 | 0 0 0 1 0 0 3 0 0 2 …  
   7 | 0 0 0 1 1 0 0 0 1 2 …  
   8 | 0 0 0 1 1 0 1 0 1 2 …  
   9 | 0 0 1 0 1 0 1 0 1 3 …  
  10 | 0 0 1 0 1 0 1 0 1 3 …  
      ...

谢谢你的帮助。

标签: pythonpandasdataframetrain-test-split

解决方案


这是我只使用 pandas 制作的一个简单函数。样本函数随机且均匀地选择数据框中的行(轴 = 0)用于测试集。可以通过删除原始数据帧中与测试集具有相同索引的行来选择训练集的行。

def train_test_split(df, frac=0.2):
    
    # get random sample 
    test = df.sample(frac=frac, axis=0)

    # get everything but the test sample
    train = df.drop(index=test.index)

    return train, test

推荐阅读