假设X,Y = load_mnist()
,其中X和Y是包含整个mnist的张量。现在我想要一小部分数据来使我的代码运行得更快,但我需要保持所有10个类,并且还要保持平衡。是否有捷径可寻?
scikit-learn的train_test_split
用于将数据拆分为训练和测试类,但您可以使用它来使用stratified
参数创建数据集的“平衡”子集。您只需指定所需的列车/测试大小比例,从而获得较小的,分层的数据样本。在你的情况下:
from sklearn.model_selection import train_test_split
X_1, X_2, Y_1, Y_2 = train_test_split(X, Y, stratify=Y, test_size=0.5)
如果您想通过更多控制来执行此操作,可以使用numpy.random.randint
生成子集大小的索引并对原始数组进行采样,如下面的代码片段所示:
# input data, assume that you've 10K samples
In [77]: total_samples = 10000
In [78]: X, Y = np.random.random_sample((total_samples, 784)), np.random.randint(0, 10, total_samples)
# out of these 10K, we want to pick only 500 samples as a subset
In [79]: subset_size = 500
# generate uniformly distributed indices, of size `subset_size`
In [80]: subset_idx = np.random.choice(total_samples, subset_size)
# simply index into the original arrays to obtain the subsets
In [81]: X_subset, Y_subset = X[subset_idx], Y[subset_idx]
In [82]: X_subset.shape, Y_subset.shape
Out[82]: ((500, 784), (500,))
X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=Ture, test_size=0.33, random_state=42)
分层将确保班级的比例。
如果你想进行K-Fold那么
from sklearn.model_selection import StratifiedShuffleSplit
sss = StratifiedShuffleSplit(n_splits=5, test_size=0.5, random_state=0)
for train_index, test_index in sss.split(X, y):
print("TRAIN:", train_index, "TEST:", test_index)
X_train, X_test = X.iloc[train_index], X.iloc[test_index]
y_train, y_test = y.iloc[train_index], y.iloc[test_index]
查看qazxsw poi获取sklearn文档。