使用Sklearn train_test_split如何指定测试集以对比预测标签与真实标签

问题解决方法 #

方案一：手动指定固定测试集（实现最后40条作为测试集的需求） #

你可以直接跳过train_test_split函数，按照索引手动切割数据集，替换原有数据划分的代码即可：

# 总样本共400条，360条为训练集，最后40条为固定测试集
split_pos = 360
Train_X, Test_X = X[:split_pos], X[split_pos:]
Train_Y, Test_Y = y[:split_pos], y[split_pos:]
# 同步保留测试集对应的原始文本与真实标签
test_raw_text = train_text[split_pos:]
test_raw_label = y[split_pos:]

该方案下测试集固定，你可以直接将预测结果和原始的40条测试数据逐行对比。

方案二：仍使用随机划分的train_test_split，保留原始索引映射 #

如果你不需要固定测试集，只是需要找到随机划分的测试集对应的原始语料条目，可以在划分时同步拆分索引数组：

# 生成全量数据的原始索引
raw_indices = np.arange(len(y))
# 拆分时同步返回训练集、测试集对应的原始索引
Train_X, Test_X, Train_Y, Test_Y, Train_indices, Test_indices = train_test_split(
    X, y, raw_indices, test_size=0.1, random_state = 1004
)
# 根据测试集索引取出对应的原始文本和真实标签
test_raw_text = [train_text[i] for i in Test_indices]
test_raw_label = [y[i] for i in Test_indices]

错误样例输出代码 #

预测完成后可以直接输出所有分类错误的条目，方便后续优化：

print("===== 分类错误条目 =====")
for text, true_lb, pred_lb in zip(test_raw_text, test_raw_label, SVM_predict_test):
    if true_lb != pred_lb:
        print(f"文本内容：{text}\n真实标签：{true_lb}，预测标签：{pred_lb}\n")

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