python-具有多个功能的train_test_split

我目前正在尝试使用决策树分类器训练数据集,但无法使train_test_split正常工作.

在下面的代码中,CS是目标输出,而功能输入是EN SN JT FT PW YR LO LA.

通过OHL的所有变量均为稀疏矩阵格式,而其他变量则为直接从数据帧获取的数组.

def OHL(x, column): #OneHotEncoder
    le = LabelEncoder()
    enc = OneHotEncoder()
    Labeled = le.fit_transform(x[column].astype(str))
    return enc.fit_transform(Labeled.reshape(-1,1))

###------------------------------------------------------------------------

df = pd.read_csv('h1b_kaggle.csv')
df = df.drop(['Unnamed: 0','WORKSITE'],1)

###------------------------------------------------------------------------

CS = OHL(df, 'CASE_STATUS')
EN = OHL(df, 'EMPLOYER_NAME')
SN = OHL(df, 'SOC_NAME')
JT = OHL(df, 'JOB_TITLE')
FT = OHL(df, 'FULL_TIME_POSITION')
PW = np.array(df['PREVAILING_WAGE'])
YR = OHL(df, 'YEAR')
LO = np.array(df['lon'])
LA = np.array(df['lat'])

解决方法:

如果查看sklearn.model_selection.train_test_split,则可以看到它带有* arrays参数.因此,要拆分前三个参数,可以使用

CS_tr, CS_te, EN_tr, EN_te, SN_tr, SN_te = train_test_split(CS, EN, SN)

(当然,您可以传递更多的数组).

请注意,当给定稀疏数组时,当前版本的sklearn将返回稀疏数组.