Python中使用iterative_train_test_split遇KeyError的解决方法

核心问题排查与修复步骤 #

1. 文本特征必须先转为数值矩阵 #

iterative_train_test_split 不支持直接传入原始文本作为特征X，必须先通过特征工程将文本转换为数值型矩阵：

• 常用工具：TfidfVectorizer（推荐）或CountVectorizer
• 代码示例：

from sklearn.feature_extraction.text import TfidfVectorizer
from imblearn.model_selection import iterative_train_test_split

# 假设数据集为df，文本列名为'text'，标签列为'label1'、'label2'
tfidf = TfidfVectorizer()
X = tfidf.fit_transform(df['text'])

2. 确保标签y的格式符合要求 #

该函数支持多标签任务，但y必须是二维数组/稀疏矩阵，不能直接传入DataFrame：

• 若标签是字符串类型，需先编码为数值；若已是数值，直接转为二维数组
• 代码示例：

from sklearn.preprocessing import LabelEncoder
import numpy as np

# 对字符串标签编码
le1 = LabelEncoder()
le2 = LabelEncoder()
y1 = le1.fit_transform(df['label1'])
y2 = le2.fit_transform(df['label2'])

# 合并为二维标签数组
y = np.column_stack((y1, y2))

3. 规范函数调用参数 #

iterative_train_test_split的标准调用格式：

X_train, X_test, y_train, y_test = iterative_train_test_split(X, y, test_size=0.2, random_state=42)

• 确认X（数值特征矩阵）与y（二维标签数组）的样本数一致（X.shape[0] == y.shape[0]）
• 避免传入非必要参数，严格遵循参数顺序

4. 常见错误快速修复 #

• 若报错Expected 2D array, got 1D array instead：检查y是否为二维，可通过y.reshape(-1, 2)强制转换
• 若报错could not convert string to float：确认X已完成文本转数值，或y的字符串标签已完成编码

完整可运行示例代码 #

import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.preprocessing import LabelEncoder
from imblearn.model_selection import iterative_train_test_split
import numpy as np

# 测试数据
data = {
    'text': ['这是测试文本1', '这是测试文本2', '这是测试文本3', '这是测试文本4', '这是测试文本5'],
    'label1': ['pos', 'neg', 'pos', 'neg', 'pos'],
    'label2': ['high', 'low', 'low', 'high', 'high']
}
df = pd.DataFrame(data)

# 文本转数值特征
tfidf = TfidfVectorizer()
X = tfidf.fit_transform(df['text'])

# 标签编码并转换为二维数组
le1 = LabelEncoder()
le2 = LabelEncoder()
y_encoded = np.column_stack((le1.fit_transform(df['label1']), le2.fit_transform(df['label2'])))

# 划分数据集
X_train, X_test, y_train, y_test = iterative_train_test_split(X, y_encoded, test_size=0.2, random_state=42)

# 验证结果
print(f"训练集样本数: {X_train.shape[0]}, 测试集样本数: {X_test.shape[0]}")
print(f"训练集标签形状: {y_train.shape}, 测试集标签形状: {y_test.shape}")

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