Python 实现训练集与测试集的随机划分

Python 实现训练集与测试集随机划分

在机器学习项目中，将数据集划分为训练集和测试集是基础且关键的步骤。以下代码展示了如何使用 Python 实现这一过程，包括读取 CSV 数据并按指定比例进行随机划分。

基础示例

首先，通过一个简单的列表示例演示随机划分的原理：

import random
dataSet = [[0], [1], [2], [3], [4], [5], [6], [7], [8], [9], [10]]
trainDataSet = random.sample(dataSet, 3)

完整实现

对于实际的 CSV 数据集，可以使用以下函数进行更复杂的处理。该函数读取 CSV 文件，跳过表头，并按指定比例随机划分训练集和测试集：

import csv
import random
def getDataSet(proportion):
  """
    :exception
      获取训练集和测试集(将数据按比例随机划分)
    :parameter
      proportion - 测试集/数据集
    :return
      trainDataSet - 训练集
      testDataSet - 测试集
    author
      肖政宇
    modify
      2019年5月10日
  """
  dataSet = open('数据集.csv')
  dataSetReader = csv.reader(dataSet)
  """
    :exception
      将数据保存到数组
  """
  dataSet = []
  next(dataSetReader, 'none') # 跳过表头
  data = next(dataSetReader, 'none')
  while (data != 'none'):
    dataSet.append(data)
    data = next(dataSetReader, 'none')
  """
    :exception
      按照比例随机划分出训练集和测试集
  """
  dataNumber = dataSet.__len__() # 数据集数据条数
  testNumber = int(dataNumber * proportion) # 测试集数据条数
  testDataSet = [] # 测试数据集
  trainDataSet = [] # 训练数据集
 
  testDataSet = random.sample(dataSet, testNumber) # 测试集
  for testData in testDataSet: # 将已经选定的测试集数据从数据集中删除
    dataSet.remove(testData)
  trainDataSet = dataSet # 训练集
 
  return trainDataSet, testDataSet

该函数通过以下步骤实现数据划分：

1. 打开 CSV 文件并创建读取器。
2. 跳过表头，将数据逐行读取到列表中。
3. 计算测试集的数据条数（基于总数据条数和指定比例）。
4. 使用 random.sample 函数随机选取测试集数据。
5. 从原始数据集中移除测试集数据，剩余部分即为训练集。
6. 返回训练集和测试集。

通过以上方法，可以灵活地将数据集划分为训练集和测试集，为后续的模型训练和评估提供基础数据。

以上就是 Python 实现训练集与测试集随机划分的详细内容，更多关于 Python 机器学习数据预处理资料请关注本站其它相关文章！