隨機切分csv訓練集和測試集

使用numpy切分訓練集和測試集

序言

在機器學習的任務中，時常需要將一個完整的數據集切分為訓練集和測試集。此處我們使用numpy完成這個任務。 iris數據集中有150條數據，我們將120條數據整合為訓練集，將30條數據整合為測試集。

iris.csv下載

程序

import csv
import os
import numpy as np

'''將iris.csv中的數據分成train_iris和test_iris兩個csv文件，其中train_iris.csv中有120個數據，test_iris.csv中有30個數據'''
labels = []
data = []
a_train_file = 'train_iris.csv'
a_test_file = 'test_iris.csv'
a_file = 'iris.csv'

seed = 3
np.random.seed(seed)
train_indices = np.random.choice(150, 120, replace=False) # 設置隨機數生成從0-150中隨機挑選120個隨機數
residue = np.array(list(set(range(150)) - set(train_indices)))
test_indices = np.random.choice(len(residue),30, replace=False) # 如果訓練集和測試集綜合的數據加起來就是一整個數據集則不需要這個操作

with open(a_file)as afile:
    a_reader = csv.reader(afile)  #從原始數據集中將所有數據讀取出來並保存到a_reader中
    labels = next(a_reader)  # 提取第一行設置為labels
    for row in a_reader:  # 將a_reader中每一行的數據提取出來並保存到data的列表中
        data.append(row)


# 生成訓練數據集
if not os.path.exists(a_train_file):
    with open(a_train_file, "w", newline='') as a_trian:
        writer = csv.writer(a_trian)
        writer.writerows([labels])  #第一行為標簽行
        writer.writerows(np.array(data)[train_indices])
        a_trian.close()

# 生成測試數據集
if not os.path.exists(a_test_file):
    with open(a_test_file, "w", newline='')as a_test:
        writer = csv.writer(a_test)
        writer.writerows([labels])  #第一行為標簽行
        writer.writerows(np.array(data)[test_indices])
        a_test.close()