Scikit-learn 中的分层培训/测试分割

我需要把我的数据分成训练集(75%)和测试集(25%)。我现在用下面的代码做到这一点:

X, Xt, userInfo, userInfo_train = sklearn.cross_validation.train_test_split(X, userInfo)

但是，我想对我的训练数据集进行分层。我该怎么做？我一直在研究 StratifiedKFold方法，但是没有指定75%/25% 的分割，只对训练数据集进行了分层。

238113 次浏览
小开
最佳答案

[update for 0.17]

请参阅 sklearn.model_selection.train_test_split的文档:

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y,
stratify=y,
test_size=0.25)

[/0.17更新]

有一个拉请求 给你。 但是你可以简单地做 train, test = next(iter(StratifiedKFold(...))) 如果你愿意，可以使用火车和测试指数。

小开

DR: 使用 分层的 ShuffleSplit和 test_size=0.25

Scikit-learn 提供两个分层分割模块:

  1. StratifiedkFold : 这个模块作为一个直接的 k- 折叠交叉验证运算符非常有用，因为它将建立 n_folds训练/测试集，这样两个类都是平衡的。

下面是一些代码(直接来自上面的文档)

>>> skf = cross_validation.StratifiedKFold(y, n_folds=2) #2-fold cross validation
>>> len(skf)
2
>>> for train_index, test_index in skf:
...    print("TRAIN:", train_index, "TEST:", test_index)
...    X_train, X_test = X[train_index], X[test_index]
...    y_train, y_test = y[train_index], y[test_index]
...    #fit and predict with X_train/test. Use accuracy metrics to check validation performance
  1. StratifiedShuffleSplit : 该模块创建一个具有同等平衡(分层)类的单一训练/测试集。本质上，这就是 n_iter=1所需要的。您可以在这里提到与 train_test_split中相同的测试大小

Code:

>>> sss = StratifiedShuffleSplit(y, n_iter=1, test_size=0.5, random_state=0)
>>> len(sss)
1
>>> for train_index, test_index in sss:
...    print("TRAIN:", train_index, "TEST:", test_index)
...    X_train, X_test = X[train_index], X[test_index]
...    y_train, y_test = y[train_index], y[test_index]
>>> # fit and predict with your classifier using the above X/y train/test
小开

下面是一个连续/回归数据示例(直到解析出 GitHub 上的这期杂志)。

min = np.amin(y)
max = np.amax(y)


# 5 bins may be too few for larger datasets.
bins     = np.linspace(start=min, stop=max, num=5)
y_binned = np.digitize(y, bins, right=True)


X_train, X_test, y_train, y_test = train_test_split(
X,
y,
stratify=y_binned
)
  • 其中 start是最小值，stop是你连续目标的最大值。
  • If you don't set right=True then it will more or less make your max value a separate bin and your split will always fail because too few samples will be in that extra bin.
小开

除了@Andreas Mueller 已经接受的回答之外，我还想补充一点，就像上面提到的@tangy:

分层 ShuffleSplit 最接近于 火车 _ 测试 _ 分离(分层 = y) with added features of:

  1. 分层 by default
  2. 通过指定 分裂，它可以重复分割数据
小开 
#train_size is 1 - tst_size - vld_size
tst_size=0.15
vld_size=0.15


X_train_test, X_valid, y_train_test, y_valid = train_test_split(df.drop(y, axis=1), df.y, test_size = vld_size, random_state=13903)


X_train_test_V=pd.DataFrame(X_train_test)
X_valid=pd.DataFrame(X_valid)


X_train, X_test, y_train, y_test = train_test_split(X_train_test, y_train_test, test_size=tst_size, random_state=13903)
小开

你可以简单的使用 Scikit 的 train_test_split()方法学习:

from sklearn.model_selection import train_test_split
train, test = train_test_split(X, test_size=0.25, stratify=X['YOUR_COLUMN_LABEL'])

I have also prepared a short GitHub Gist which shows how stratify option works:

https://gist.github.com/SHi-ON/63839f3a3647051a180cb03af0f7d0d9

小开

从上面更新@tangy 的答案到 scikit-learn 的当前版本: 0.23.2(StratifiedShuffleSplit documentation)。

from sklearn.model_selection import StratifiedShuffleSplit


n_splits = 1  # We only want a single split in this case
sss = StratifiedShuffleSplit(n_splits=n_splits, test_size=0.25, random_state=0)


for train_index, test_index in sss.split(X, y):
X_train, X_test = X[train_index], X[test_index]
y_train, y_test = y[train_index], y[test_index]
小开

因此，可取的做法是将数据集划分为训练集和测试集，以便在每个类中保留与在原始数据集中观察到的相同比例的示例。

这就是所谓的分层列车试验分裂。

我们可以通过将“分层”参数设置为原始数据集的 y 组件来实现这一点。Train _ test _ split ()函数将使用这个函数来确保 train 和测试集在所提供的“ y”数组中的每个类中具有示例的比例。

小开

StratifiedShuffleSplit 是在我们选择了应该在即将生成的所有小数据集中均匀表示的列之后完成的。 这些褶皱是通过保留每个类别样品的百分比而制成的

假设我们有一个数据集“ data”，其中有一列是“ Season”，我们希望得到一个“ Season”的均衡表示，那么它看起来是这样的:

from sklearn.model_selection import StratifiedShuffleSplit
sss=StratifiedShuffleSplit(n_splits=1,test_size=0.25,random_state=0)


for train_index, test_index in sss.split(data, data["season"]):
sss_train = data.iloc[train_index]
sss_test = data.iloc[test_index]

提交答案