python - 训练测试拆分以确保所有类别都包含在训练集中

Question

假设数据中有大约 20 个分类列，每个都有一组不同的唯一分类值。现在必须进行训练测试拆分，并且需要确保所有唯一类别都包含在训练集中。怎么做到呢？我还没有尝试过，但是所有这些列都应该包含在分层参数中吗？

标签： pythoncategorical-datatrain-test-split

Answer 1

是的。这是正确的。

为了演示，我正在使用Melbourne Housing Dataset。

import pandas as pd
from sklearn.model_selection import train_test_split

Meta = pd.read_csv('melb_data.csv')
Meta = Meta[["Rooms", "Type", "Method", "Bathroom"]]
print(Meta.head())

print("\nBefore split -- Method feature distribution\n")
print(Meta.Method.value_counts(normalize=True))
print("\nBefore split -- Type feature distribution\n")
print(Meta.Type.value_counts(normalize=True))

train, test = train_test_split(Meta, test_size = 0.2, stratify=Meta[["Method", "Type"]])

print("\nAfter split -- Method feature distribution\n")
print(train.Method.value_counts(normalize=True))
print("\nAfter split -- Type feature distribution\n")
print(train.Type.value_counts(normalize=True))

输出

Rooms Type Method  Bathroom
0      2    h      S       1.0
1      2    h      S       1.0
2      3    h     SP       2.0
3      3    h     PI       2.0
4      4    h     VB       1.0

Before split -- Method feature distribution

S     0.664359
SP    0.125405
PI    0.115169
VB    0.088292
SA    0.006775
Name: Method, dtype: float64

Before split -- Type feature distribution

h    0.695803
u    0.222165
t    0.082032
Name: Type, dtype: float64

After split -- Method feature distribution

S     0.664396
SP    0.125368
PI    0.115151
VB    0.088273
SA    0.006811
Name: Method, dtype: float64

After split -- Type feature distribution

h    0.695784
u    0.222202
t    0.082014
Name: Type, dtype: float64