Cómo dividir datos en un conjunto de entrenamiento equilibrado y un conjunto de prueba en sklearn

Solución:

Aunque la sugerencia de Christian es correcta, técnicamente train_test_split debería darle resultados estratificados mediante el uso de stratify param.

Entonces podrías hacer:

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

El truco aquí es que comienza desde la versión 0.17 en sklearn.

De la documentación sobre el parámetro stratify:

estratificar: similar a una matriz o Ninguno (el valor predeterminado es Ninguno) Si no es Ninguno, los datos se dividen de forma estratificada, utilizando esto como la matriz de etiquetas. Nuevo en la versión 0.17: estratificar la división

Puede usar StratifiedShuffleSplit para crear conjuntos de datos con el mismo porcentaje de clases que el original:

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]

Si las clases no están equilibradas pero desea que la división esté equilibrada, la estratificación no ayudará. No parece haber un método para hacer un muestreo equilibrado en sklearn, pero es bastante fácil usar numpy básico, por ejemplo, una función como esta podría ayudarlo:

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

Tenga en cuenta que si usa esto y muestrea más puntos por clase que en los datos de entrada, entonces esos serán muestreados (muestra con reemplazo). Como resultado, algunos puntos de datos aparecerán varias veces y esto puede tener un efecto en las medidas de precisión, etc. Y si alguna clase tiene solo un punto de datos, habrá un error. Puede comprobar fácilmente el número de puntos por clase, por ejemplo, con np.unique(target, return_counts=True)

---

---

Utiliza Nuestro Buscador