将训练和测试数据保存到文件中

Question

我正在使用以下代码将数据集分成训练数据和测试数据，并保存到一个文件中；

import numpy as np
from sklearn.cross_validation import train_test_split

a = (np.genfromtxt(open('dataset.csv','r'), delimiter=',', dtype='int')[1:])
a_train, a_test = train_test_split(a, test_size=0.33, random_state=0)

c1 = open('trainfile.csv', 'w')
arr1 = str(a_train)
c1.write(arr1)
c1.close

c2 = open('testfile.csv', 'w')
arr2 = str(a_test)
c2.write(arr2)
c2.close

但是我在文件中得到了以下输出；

trainfile.csv:
[[ 675847       0       0 ...,       0       0       3]
 [  74937       0       0 ...,       0       0       3]
 [  65212       0       0 ...,       0       0       3]
 ..., 
 [  18251       0       0 ...,       0       0       1]
 [1131828       0       0 ...,       0       0       1]
 [  14529       0       0 ...,       0       0       1]]

这就是trainfile的全部内容。我在testfile.csv的输出上也遇到了同样的问题。我想要的是将完整的训练数据和测试数据存储在文件中，而不是用点表示额外的数据。有什么建议吗？

Answer 1

这是因为你在numpy数组上调用了字符串方法str。应该使用numpy提供的函数numpy.savetxt。用法大概是这样的：

with open('testfile.csv', 'w') as FOUT:
    np.savetxt(FOUT, a_test)

需要注意的是，这种格式不一定能被CSV阅读器识别。如果你想要这样的效果，可以使用https://docs.python.org/2/library/csv.html。