如何在sklearn上将数据分割为平衡的训练集和测试集

48

我在多分类任务中使用sklearn。我需要将所有数据拆分为训练集和测试集。我希望从每个类别中随机抽取相同数量的样本。 实际上,我正在使用这个函数

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

但是这会导致数据集不平衡!有什么建议。


如果您仍然想使用cross_validation.train_test_split,并且您正在使用sklearn 0.17,则可以平衡训练和测试,请查看我的答案。 - Guiem Bosch
1
顺便提一下,对于一个不平衡的训练集,例如使用sklearn.ensemble.RandomForestClassifier,可以使用class_weight="balanced" - Shadi
@Shadi:请注意,平衡您的训练集是一件不同的事情;class_weight将对您的成本最小化产生影响。 - Markus
这个函数对我来说似乎不太有趣 ;) - Kartik Chugh
6个回答

48

虽然Christian的建议是正确的,但从技术上讲,train_test_split应该使用stratify参数给出分层结果。

所以您可以这样做:

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

这里的技巧是它从版本0.17开始在sklearn中。

从有关参数stratify的文档中:

stratify:类似数组或None(默认值为None) 如果不为None,则使用此数组作为标签以分层方式拆分数据。 新版本0.17中:分层拆分


5
但是如果数据中的类别不平衡(例如:class1=200个样本,class2=250个样本...),我需要从中选取(100, 100)进行训练,(50, 50)进行测试。我该如何做? - Jeanne
1
train_test_split 中有两个额外的参数:train_sizetest_size(除了表示比例之外,它们也可以是 int)。虽然我从未尝试过,但我认为 train_size=100test_size=50 结合 stratify 参数应该可以工作。 - Guiem Bosch
2
我没有尝试过,但如果你这样做,你应该有100个遵循原始分布的训练样本和50个也遵循原始分布的训练样本。(我将稍微改变一下例子以澄清,假设class1=200个样本,class2=400个样本),那么你的训练集将有33个来自class1和67个来自class2,而你的测试集将有18个来自class1和32个来自class2。 据我所知,原问题是想要一个由50个class1示例和50个class2示例组成的训练集,但测试集中有18个class1示例和32个class2示例。 - Rodrigo Laguna
3
澄清一下,使用分层抽样(split using stratify)可按照原始数据的相同比例创建数据样本。例如,如果您的数据中的类别被划分为70/30,则分层抽样会创建70/30的样本。 - BenP

36
你可以使用StratifiedShuffleSplit来创建具有与原始数据集相同类别百分比的数据集:
import numpy as np
from sklearn.model_selection import StratifiedShuffleSplit
X = 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]

8
注意:自版本0.18起,StratifiedShuffleSplit已被弃用:该模块将在0.20中被删除。请改用sklearn.model_selection.StratifiedShuffleSplit - mc2
创建数据集,其中包含与原始数据集相同比例的类别:”根据https://github.com/scikit-learn/scikit-learn/issues/8913,这并不总是成立。 - gented
代码未经测试,我猜测是因为出现了“stratSplit不可迭代”的错误。 - Pfinnn

12

如果类别不平衡但你想要分割是均衡的,那么分层并不能帮助你。在sklearn中似乎没有一种方法来进行平衡抽样,但使用基本的numpy很容易实现,例如下面这个函数可能会对你有所帮助:

def split_balanced(data, target, test_size=0.2):

    classes = np.unique(target)
    # can give test_size as fraction of input data size of number of samples
    if test_size<1:
        n_test = np.round(len(target)*test_size)
    else:
        n_test = test_size
    n_train = max(0,len(target)-n_test)
    n_train_per_class = max(1,int(np.floor(n_train/len(classes))))
    n_test_per_class = max(1,int(np.floor(n_test/len(classes))))

    ixs = []
    for cl in classes:
        if (n_train_per_class+n_test_per_class) > np.sum(target==cl):
            # if data has too few samples for this class, do upsampling
            # split the data to training and testing before sampling so data points won't be
            #  shared among training and test data
            splitix = int(np.ceil(n_train_per_class/(n_train_per_class+n_test_per_class)*np.sum(target==cl)))
            ixs.append(np.r_[np.random.choice(np.nonzero(target==cl)[0][:splitix], n_train_per_class),
                np.random.choice(np.nonzero(target==cl)[0][splitix:], n_test_per_class)])
        else:
            ixs.append(np.random.choice(np.nonzero(target==cl)[0], n_train_per_class+n_test_per_class,
                replace=False))

    # take same num of samples from all classes
    ix_train = np.concatenate([x[:n_train_per_class] for x in ixs])
    ix_test = np.concatenate([x[n_train_per_class:(n_train_per_class+n_test_per_class)] for x in ixs])

    X_train = data[ix_train,:]
    X_test = data[ix_test,:]
    y_train = target[ix_train]
    y_test = target[ix_test]

    return X_train, X_test, y_train, y_test

请注意,如果您使用此方法并对每个类别采样的点数多于输入数据,则这些点将被上采样(带有替换的采样)。因此,某些数据点会出现多次,这可能会影响准确度等指标。如果某个类别只有一个数据点,则会出现错误。您可以轻松地检查每个类别的数据点数量,例如使用np.unique(target, return_counts=True)

1
我喜欢这个原则,但是我认为当前实现存在问题,即随机抽样可能会将相同的样本分配给训练集和测试集。采样应该从不同的池中收集训练和测试索引。 - DonSteep
2
你说得完全正确,我试图通过说“你的训练和测试数据中可能存在重复点,这可能会导致模型性能看起来过于乐观”来提及这一点,但我现在明白措辞可能不够完美,对此很抱歉。我将编辑代码,以便不再有共享数据点。 - antike
1
我不确定你的帖子是否准确。当你提到“平衡”时,是指每个类别的比例大约相等吗?还是指测试集中的类别分布与训练集中的类别分布大致相同。分层抽样可以实现后者。 - JoAnn Alvarez

1
另一种方法是从分层测试/训练拆分中过度或欠采样。这时可以使用imbalanced-learn库,非常方便,特别适用于进行在线学习并希望在管道内保证平衡的训练数据。
from imblearn.pipeline import Pipeline as ImbalancePipeline

model = ImbalancePipeline(steps=[
  ('data_balancer', RandomOverSampler()),
  ('classifier', SVC()),
])

1
这是我正在使用的函数。你可以适应它并进行优化。
# Returns a Test dataset that contains an equal amounts of each class
# y should contain only two classes 0 and 1
def TrainSplitEqualBinary(X, y, samples_n): #samples_n per class
    
    indicesClass1 = []
    indicesClass2 = []
    
    for i in range(0, len(y)):
        if y[i] == 0 and len(indicesClass1) < samples_n:
            indicesClass1.append(i)
        elif y[i] == 1 and len(indicesClass2) < samples_n:
            indicesClass2.append(i)
            
        if len(indicesClass1) == samples_n and len(indicesClass2) == samples_n:
            break
    
    X_test_class1 = X[indicesClass1]
    X_test_class2 = X[indicesClass2]
    
    X_test = np.concatenate((X_test_class1,X_test_class2), axis=0)
    
    #remove x_test from X
    X_train = np.delete(X, indicesClass1 + indicesClass2, axis=0)
    
    Y_test_class1 = y[indicesClass1]
    Y_test_class2 = y[indicesClass2]
    
    y_test = np.concatenate((Y_test_class1,Y_test_class2), axis=0)
    
    #remove y_test from y
    y_train = np.delete(y, indicesClass1 + indicesClass2, axis=0)
    
    if (X_test.shape[0] != 2 * samples_n or y_test.shape[0] != 2 * samples_n):
        raise Exception("Problem with split 1!")
        
    if (X_train.shape[0] + X_test.shape[0] != X.shape[0] or y_train.shape[0] + y_test.shape[0] != y.shape[0]):
        raise Exception("Problem with split 2!")
    
    return X_train, X_test, y_train, y_test

0

这是我用来获取训练/测试数据索引的实现

def get_safe_balanced_split(target, trainSize=0.8, getTestIndexes=True, shuffle=False, seed=None):
    classes, counts = np.unique(target, return_counts=True)
    nPerClass = float(len(target))*float(trainSize)/float(len(classes))
    if nPerClass > np.min(counts):
        print("Insufficient data to produce a balanced training data split.")
        print("Classes found %s"%classes)
        print("Classes count %s"%counts)
        ts = float(trainSize*np.min(counts)*len(classes)) / float(len(target))
        print("trainSize is reset from %s to %s"%(trainSize, ts))
        trainSize = ts
        nPerClass = float(len(target))*float(trainSize)/float(len(classes))
    # get number of classes
    nPerClass = int(nPerClass)
    print("Data splitting on %i classes and returning %i per class"%(len(classes),nPerClass ))
    # get indexes
    trainIndexes = []
    for c in classes:
        if seed is not None:
            np.random.seed(seed)
        cIdxs = np.where(target==c)[0]
        cIdxs = np.random.choice(cIdxs, nPerClass, replace=False)
        trainIndexes.extend(cIdxs)
    # get test indexes
    testIndexes = None
    if getTestIndexes:
        testIndexes = list(set(range(len(target))) - set(trainIndexes))
    # shuffle
    if shuffle:
        trainIndexes = random.shuffle(trainIndexes)
        if testIndexes is not None:
            testIndexes = random.shuffle(testIndexes)
    # return indexes
    return trainIndexes, testIndexes

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接