Python: 如何将数据分为测试集和训练集？

Question

Python: 如何将数据分为测试集和训练集？

3

我一直在使用CSV数据来实现我的脚本，并希望将数据分成两个数据集进行抽样：

测试数据
训练数据

我想将数据集按85%和15%的比例进行抽样，并输出两个CSV文件Test.csv和Train.csv。

我希望这可以使用基本的Python来实现，而不需要使用任何其他外部模块，如NumPy、SciPy、Pandas或Scikitlearn。有人能帮我按百分比随机抽样数据吗？此外，我将提供包含随机观察数量的数据集。到目前为止，我只了解了Pandas和其他各种模块来按百分比抽样数据，并没有得到针对我的问题的具体解决方案。

此外，我希望在两个文件中保留CSV头。因为标题可以使每行可访问，并可用于进一步分析。

- desmond.carros

你的帖子太笼统了。详细说明你已经尝试过什么。使用问号明确表达你所提出的问题。 - Martin Cowie

@MartinCowie 我刚刚在网上学习了一下搜索。到目前为止还没有尝试过任何东西。我正在寻找一种逻辑，想要从现有文件中创建两个文件。Test.csv和Train.csv来自主文件data.csv。我希望85%的数据应该在test.csv中，其余15%的数据在train.csv中。 - desmond.carros

为什么您想将85%的数据作为测试数据，15%作为训练数据？很可能是因为您需要85%的数据进行训练，剩下的作为测试数据。 - Anup Verma

2个回答

2

使用随机模块中的random函数获取在0和1之间均匀分布的随机数。

如果它大于.85，则写入训练数据，否则写入测试数据。请参见如何在Python中模拟有偏硬币的翻转？。

import random

with open(input_file) as data:
    with open(test_output, 'w') as test:
        with open(train_output, 'w') as train:
            header = next(data)
            test.write(header)
            train.write(header)
            for line in data:
                if random.random() > 0.85:
                    train.write(line)
                else:
                    test.write(line)

- Peter Wood

csv.writer没有名为'write'的对象，我认为应该是writerows？ - desmond.carros

向你致敬！但数据也包含了标题。而且这些标题在一个文件中，而不在第二个文件中。有没有办法在两个文件中都保留这些标题？ - desmond.carros

1

@desmond.carros：在这种情况下，csv.DictReader和csv.DictWriter可能会很有用。 - Callidior

我也是这么想的。 :) - desmond.carros

1

@desmond.carros 把那个放在问题里，或者提出/搜索另一个问题。 - Peter Wood

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Callidior · Accepted Answer

使用random.shuffle来创建数据集的随机排列，并按照您的需求进行切片：

import random
random.shuffle(data)
train = data[:int(len(data)*0.85)]
test = data[len(train):]

既然您要求一种特定的解决方案来将可能很大的CSV文件分成两个文件，以用于训练和测试数据，我也会展示如何使用与上述通用方法类似的方法来完成这项任务：

import random

# Count lines
with open('data.csv','r') as csvf:
    linecount = sum(1 for lines in csvf if line.strip() != '')

# Create index sets for training and test data
indices = list(range(linecount))
random.shuffle(indices)
ind_test = set(indices[:int(linecount*0.15)])
del indices

# Partition CSV file
with open('data.csv','r') as csvf, open('train.csv','w') as trainf, open('test.csv','w') as testf:
    i = 0
    for line in csvf:
        if line.strip() != '':
            if i in ind_test:
                testf.write(line.strip() + '\n')
            else:
                trainf.write(line.strip() + '\n')

因此，我认为CSV文件每行包含一个观察结果。

这将创建一个准确的85:15分割。如果对您来说不需要非常准确的分区，则Peter Wood的解决方案会更加高效。