Scikit-Learn 提供了一些函数，可以用多种方式将数据集分割成多个子集。
sklearn.model_selection.train_test_split
是纯随机的取样方法，即没有对原数据集进行分层，具体调用如下：
from sklearn.model_selection import train_test_split
train_set, test_set = train_test_split(data, test_size=0.2, random_state=42)
其中，random_state=42是设置随机生成器种子，具体原理可参考【sklearn.model_selection.train_test_split(random_state=42) 与 np.random.seed(42)的联系和原理详解】
sklearn.model_selection.StratifiedShuffleSplit
你的数据集很大时(尤其是和属性数相比)，纯随机的取样方法通常可行;但如果数据集不大，就会有采样偏差的风险。当一个调查公司想要对 1000 个人进行调查，它们不是在电话亭里随机选 1000 个人出来。调查公司要保证这 1000 个人对人群整体有代表性。例如，美国人口的 51.3% 是女性，48.7% 是男性。所以在美国， 严谨的调查需要保证样本也是这个比例:513 名女性，487 名男性。这称作分层采样 (stratified sampling):将人群分成均匀的子分组，称为分层，从每个分层去取合适数量的实例，以保证测试集对总人数有代表性。如果调查公司采用纯随机采样，会有 12% 的概率导致采样偏差:女性人数少于 49%，或多于 54%。不管发生那种情况，调查结果都会严重偏差。
假设专家告诉你，收入中位数是预测房价中位数非常重要的属性。你可能想要保证测试集可 以代表整体数据集中的多种收入分类。因为收入中位数是一个连续的数值属性，你首先需要 创建一个收入类别属性。再仔细地看一下收入中位数的柱状图。
大多数的收入中位数的值聚集在 2-5(万美元)，但是一些收入中位数会超过 6。数据集中的每个分层都要有足够的实例位于你的数据中，这点很重要。否则，对分层重要性的评估就会有偏差。这意味着，你不能有过多的分层，且每个分层都要足够大。
后面的代码通过将收入 中位数除以 1.5(以限制收入分类的数量)，创建了一个收入类别属性，用 ceil 对值舍入 (以产生离散的分类)，然后将所有大于5的分类归入到分类 5:
housing["income_cat"] = np.ceil(housing["median_income"] / 1.5) housing["income_cat"].where(housing["income_cat"] < 5, 5.0, inplace=True)
现在，就可以根据收入分类，进行分层采样。你可以使用 Scikit-Learn 的 StratifiedShuffleSplit 类:
from sklearn.model_selection import StratifiedShuffleSplit
split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
for train_index, test_index in split.split(housing, housing["income_cat"]):
strat_train_set = housing.loc[train_index]
strat_test_set = housing.loc[test_index] # 保证测试集
检查下结果是否符合预期。你可以在完整的房产数据集中查看收入分类比例:
strat_test_set["income_cat"].value_counts() / len(strat_test_set)
3 0.350533
2 0.318798
4 0.176357
5 0.114583
1 0.039729
Name: income_cat, dtype: float64
housing["income_cat"].value_counts() / len(housing)
3.0 0.350581
2.0 0.318847
4.0 0.176308
5.0 0.114438
1.0 0.039826
Name: income_cat, dtype: float64
可以看出测试集中的收入分类比例和总收入集的收入分类比例几乎相同。
现在，需要删除 income_cat 属性，使数据回到初始状态:
for set_ in (strat_train_set, strat_test_set):
set_.drop("income_cat", axis=1, inplace=True)
我们用了大量时间来生成测试集的原因是:测试集通常被忽略，但实际是机器学习非常重要的一部分。
两者区别
分层抽样是全局特征测试，保证统一分布的训练集和测试集，使得全局拟合的效果达到最好，但是会放弃一部分局部特征，导致局部特征拟合效果不是很好。因为如果测试集，训练集都是基于同一分布的，那其实训练出来的结果很大概率会是过拟合的。而且，分层抽样导致部分样本被抽中的概率变大或者变小的，导致数据不是那种均等被抽中的随机性。纯随机采样得到的测试集，可以理解为可能是局部特征测试，也可能是全局特征测试。
参考：
《Hands-On Machine Learning with Scikit-Learn and Teansorflow》