请考虑以下代码:
import pandas as pd
from sklearn.model_selection import train_test_split
# step 1
ids = list(range(1000))
label = 500 * [1.0] + 500 * [0.0]
df = pd.DataFrame({"id": ids, "label": label})
# step 2
train_p = 0.8
val_p = 0.1
test_p = 0.1
# step 3
n_train = int(len(df) * train_p)
n_val = int(len(df) * val_p)
n_test = len(df) - n_train - n_val
print("* Step 3")
print("train:", n_train)
print("val:", n_val)
print("test:", n_test)
print()
# step 4
train_ids, test_ids = train_test_split(df["id"], stratify=df.label, test_size=n_test, random_state=42)
# step 5
print("* Step 5. First split")
print( df.loc[df.id.isin(train_ids), "label"].value_counts() )
print( df.loc[df.id.isin(test_ids), "label"].value_counts() )
print()
# step 6
train_ids, val_ids = train_test_split(train_ids, stratify=df.loc[df.id.isin(train_ids), "label"], test_size=n_val, random_state=42)
# step 7
train_df = df[df["id"].isin(train_ids)]
val_df = df[df["id"].isin(val_ids)]
test_df = df[df["id"].isin(test_ids)]
# step 8
print("* Step 8. Final split")
print("train:", train_df["label"].value_counts())
print("val:", val_df["label"].value_counts())
print("test:", test_df["label"].value_counts())
输出:
* Step 3
train: 800
val: 100
test: 100
* Step 5. First split
label
1.0 450
0.0 450
Name: count, dtype: int64
label
1.0 50
0.0 50
Name: count, dtype: int64
* Step 8. Final split
train: label
0.0 404
1.0 396
Name: count, dtype: int64
val: label
1.0 54
0.0 46
Name: count, dtype: int64
test: label
1.0 50
0.0 50
Name: count, dtype: int64
- 创建一个包含 1000 个元素的数据框,这些元素在类别 1 和类别 0(正和负)之间完美平衡;
- 定义训练、验证和测试分区中样本的比例。我希望训练分区包含 800 个样本,其他两个分区各包含 100 个样本。
- 计算三个分区的大小并打印它们的值。
- 执行第一次拆分以获取测试集,分层
label
。 - 打印第一个分割的标签统计信息。两个分区仍然是平衡的。
- 执行第二次拆分,分为训练和验证,分层
label
。 - 选择示例
- 打印标签统计数据。
如您所见,步骤 6 中的第二次拆分并未产生平衡的拆分(统计数据打印于步骤 8)。第一次拆分后,示例(步骤 5 的输出)仍然是平衡的,因此可以进行第二次拆分以保持完美的类别平衡。
我做错了什么?
您只在第二个拆分中提供了 ID,这可能无法正确关联到整行数据(包含标签列),从而无法进行正确的分层。我认为它无法按预期工作,因为标签分布并未在整行中得到维护。
例如,在第一次分割中,您有一个数据框并使用了 df["id"] 和 df.label,如下所示。
因此,在第二次拆分时,您可以执行相同的操作,如下所示:
它会完美地工作!