python 数据拆分成训练集和测试集

import numpy as np
from sklearn.model_selection import train_test_split
my_matrix = np.loadtxt(open("xxxxx.csv"),delimiter=",",skiprows=0)
X, y = my_matrix[:,:-1],my_matrix[:,-1]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
train= np.column_stack((X_train,y_train))
np.savetxt('train_usual.csv',train, delimiter = ',')
test = np.column_stack((X_test, y_test))
np.savetxt('test_usual.csv', test, delimiter = ',')

 

运行程序，会发现报错：这个代码有缺陷，会把csv中的float类型读成str类型。

新建文件夹，运行如下代码：

import numpy as np

fr = open("sh.csv", 'r')  #原始数据集
frw = open("sh1.csv", 'w')

line = fr.readlines()
for L in line:
    string = L.strip("\n").split(",")
    a = np.float(string[0])
    b = np.float(string[1])
    c = np.float(string[2])
    d = np.float(string[3])
    e = np.float(string[4])
    f = np.float(string[5])
    g = np.float(string[6])

    str = '%f,%f,%f,%f,%f,%f,%f\n' % (a, b, c, d, e, f, g)  #根据你自己的数据集修改
    print(str)
    frw.write(str)
frw.close()
fr.close()


再将开始的代码运行一遍：

import numpy as np
from sklearn.model_selection import train_test_split
my_matrix = np.loadtxt(open("sh1.csv"),delimiter=",",skiprows=0)
X, y = my_matrix[:,:-1],my_matrix[:,-1]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
train= np.column_stack((X_train,y_train))
np.savetxt('train_usual.csv',train, delimiter = ',')
test = np.column_stack((X_test, y_test))
np.savetxt('test_usual.csv', test, delimiter = ',')

完成！！！