Parámetro "estratificar" del método "train_test_split" (scikit Learn)

Estoy intentando usar el train_test_splitpaquete scikit Learn, pero tengo problemas con el parámetro stratify. A continuación está el código:

from sklearn import cross_validation, datasets 

X = iris.data[:,:2]
y = iris.target

cross_validation.train_test_split(X,y,stratify=y)

Sin embargo, sigo teniendo el siguiente problema:

raise TypeError("Invalid parameters passed: %s" % str(options))
TypeError: Invalid parameters passed: {'stratify': array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2])}

¿Alguien tiene una idea de lo que está pasando? A continuación se muestra la documentación de la función.

[...]

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 matriz de etiquetas.

Nuevo en la versión 0.17: estratificar la división

[...]

Scikit-Learn solo le dice que no reconoce el argumento "estratificar", no que lo esté usando incorrectamente. Esto se debe a que el parámetro se agregó en la versión 0.17 como se indica en la documentación que citó.

Así que solo necesitas actualizar Scikit-Learn.

Este stratifyparámetro realiza una división para que la proporción de valores en la muestra producida sea la misma que la proporción de valores proporcionados al parámetro stratify.

Por ejemplo, si la variable yes una variable categórica binaria con valores 0y 1hay un 25% de ceros y un 75% de unos, stratify=yse asegurará de que su división aleatoria tenga un 25% de 0y un 75% de 1.

Para mi yo futuro que viene aquí a través de Google:

train_test_splitestá ahora en model_selection, por lo tanto:

from sklearn.model_selection import train_test_split

# given:
# features: xs
# ground truth: ys

x_train, x_test, y_train, y_test = train_test_split(xs, ys,
                                                    test_size=0.33,
                                                    random_state=0,
                                                    stratify=ys)

es la forma de utilizarlo. Establecer el random_statees deseable para la reproducibilidad.

En este contexto, la estratificación significa que el método train_test_split devuelve subconjuntos de entrenamiento y prueba que tienen las mismas proporciones de etiquetas de clase que el conjunto de datos de entrada.

Intente ejecutar este código, "simplemente funciona":

from sklearn import cross_validation, datasets 

iris = datasets.load_iris()

X = iris.data[:,:2]
y = iris.target

x_train, x_test, y_train, y_test = cross_validation.train_test_split(X,y,train_size=.8, stratify=y)

y_test

array([0, 0, 0, 0, 2, 2, 1, 0, 1, 2, 2, 0, 0, 1, 0, 1, 1, 2, 1, 2, 0, 2, 2,
       1, 2, 1, 1, 0, 2, 1])