我在使用sklearn进行多分类任务。我需要将所有数据(alldata)分割成训练集(train_set)和测试集(test_set)。我想从每个类别中随机抽取相同数量的样本。实际上，我正在使用这个函数

X_train, X_test, y_train, y_test = cross_validation.train_test_split(Data, Target, test_size=0.3, random_state=0)

但是它产生了不均衡的数据集！有什么建议吗？

回答：

你可以使用StratifiedShuffleSplit来创建与原始数据集相同比例的类别数据集：

import numpy as npfrom sklearn.model_selection import StratifiedShuffleSplitX = np.array([[1, 3], [3, 7], [2, 4], [4, 8]])y = np.array([0, 1, 0, 1])stratSplit = StratifiedShuffleSplit(y, n_iter=1, test_size=0.5, random_state=42)for train_idx, test_idx in stratSplit:    X_train=X[train_idx]    y_train=y[train_idx]print(X_train)# [[3 7]#  [2 4]]print(y_train)# [1 0]