Parâmetro “estratificar” do método “train_test_split” (scikit Learn)

Estou tentando usar train_test_splitdo pacote scikit Learn, mas estou tendo problemas com o parâmetro stratify. A seguir está o código:

from sklearn import cross_validation, datasets 

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

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

No entanto, continuo tendo o seguinte 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])}

Alguém tem ideia do que está acontecendo? Abaixo está a documentação da função.

[...]

stratify : array-like ou None (o padrão é None)

Se não for nenhum, os dados são divididos de forma estratificada, usando isso como a matriz de rótulos.

Novo na versão 0.17: divisão estratificada

[...]

O Scikit-Learn está apenas dizendo a você que não reconhece o argumento "estratificar", não que você o esteja usando incorretamente. Isso ocorre porque o parâmetro foi adicionado na versão 0.17 conforme indicado na documentação que você citou.

Então você só precisa atualizar o Scikit-Learn.

Este stratifyparâmetro faz uma divisão de forma que a proporção dos valores na amostra produzida seja a mesma que a proporção dos valores fornecidos ao parâmetro stratify.

Por exemplo, se a variável yé uma variável categórica binário com valores 0e 1e há 25% de zeros e 75% dos queridos, stratify=yvai se certificar de que a sua divisão aleatória tem 25% de 0's e 75% de 1' s.

Para meu futuro eu, que vem aqui via Google:

train_test_splitestá agora dentro model_selection, portanto:

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)

é a maneira de usá-lo. A configuração de random_stateé desejável para reprodutibilidade.

Nesse contexto, a estratificação significa que o método train_test_split retorna subconjuntos de treinamento e teste que têm as mesmas proporções de rótulos de classe que o conjunto de dados de entrada.

Tente executar este código, ele "simplesmente 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])