如何使用sklearn按writer列拆分机器学习训练测试集？

你这个需求属于按固定分组规则拆分数据集，不需要使用sklearn中做随机拆分的train_test_split，直接基于pandas的布尔索引筛选即可，输出结果可以直接适配sklearn全流程的接口要求。

核心实现代码 #

假设你存储数据的DataFrame变量名为df，存储文本内容的列名为text，机器学习任务的预测标签列名为label，可以按如下逻辑拆分：

# 定义归属训练集的作者集合
train_author_set = {"Aeschylus", "Sophocles"}

# 筛选训练集：writer列取值在训练作者集合内的样本
train_data = df[df["writer"].isin(train_author_set)].reset_index(drop=True)
# 筛选测试集：writer列取值为Euripides的样本
test_data = df[df["writer"] == "Euripides"].reset_index(drop=True)

# 拆分特征、标签，格式可直接传入sklearn模型
X_train, y_train = train_data["text"], train_data["label"]
X_test, y_test = test_data["text"], test_data["label"]

注意事项 #

• 这种拆分方式实现了作者维度的严格隔离，不会出现同一作者的样本同时出现在训练集、测试集的情况，能有效避免数据泄露，比随机拆分更贴合跨作者泛化性验证的需求。
• 如果后续需要做交叉验证，不要使用默认的随机K折拆分，要使用sklearn内置的GroupKFold，传入writer列作为分组依据，保证每一折的验证集作者都未在对应训练集中出现，示例代码如下：

from sklearn.model_selection import GroupKFold

# 初始化分组5折交叉验证
gkf = GroupKFold(n_splits=5)
# 拆分时传入groups参数指定分组列
for train_idx, val_idx in gkf.split(X=df["text"], y=df["label"], groups=df["writer"]):
    X_fold_train, X_fold_val = df["text"].iloc[train_idx], df["text"].iloc[val_idx]
    y_fold_train, y_fold_val = df["label"].iloc[train_idx], df["label"].iloc[val_idx]
    # 写入单折模型训练、验证的逻辑即可

• 筛选得到的X_train、X_test、y_train、y_test格式完全兼容sklearn的Pipeline、模型训练、评估等全流程接口，不需要额外做格式转换。

内容的提问来源于stack exchange，提问作者mar_k