主要是基于sklearn的make_regression与make_classification生成数据
import pandas as pd
def load_data(samples=1000, classification: bool = True, split_train_size: float = 0.0):
"""
用来生成训练、测试数据
:param samples: 数据量
:param classification: True为分类任务，False为回归任务
:param split_train_size: 训练集的大小 <= 1，若为0.8，则训练集80%，测试集20%
:return: 返回x与y或切分训练集后的x与y
"""
from sklearn.datasets import make_regression, make_classification
from sklearn.model_selection import train_test_split
if classification: # 生成分类模型
data_x, data_y = make_classification(n_samples=samples, n_classes=4, n_features=10, n_informative=8)
else: # 生成回归
data_x, data_y = make_regression(n_samples=samples, n_features=10)
df_x = pd.Dataframe(data_x, columns=['f_1', 'f_2', 'f_3', 'f_4', 'f_5', 'f_6', "f_7", "f_8", "f_9", "f_10"])
df_y = pd.Series(data_y)
if split_train_size:
x_train, x_test, y_train, y_test = train_test_split(df_x, df_y, train_size=0.8, random_state=0, shuffle=True)
return x_train, y_train, x_test, y_test
else:
return df_x, df_y
if __name__ == '__main__':
x, y = load_data(samples=10000, classification=False) # 回归任务数据
train_x, test_x, train_y, test_y = load_data(split_train_size=0.8) # 分类任务(默认)，同时切分训练集测试集
