填充缺失值、生成哑变量后训练与测试集列数不一致的问题解决

问题场景 #

处理含缺失值的分类变量时执行了以下数据预处理操作：

1. 拆分数据集 #

random_state_value = 0

# Define target
X = data.drop(columns='income', axis=1)
y = data['income']

# Split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=random_state_value)

2. 缺失值填充 #

# Impute missing data
imputer_cat = SimpleImputer(strategy='most_frequent')
imputer_num = SimpleImputer(strategy='median')

X_train[['workclass', 'occupation', 'native-country']] = imputer_cat.fit_transform(X_train[['workclass', 'occupation', 'native-country']])
X_train[['age']] = imputer_num.fit_transform(X_train[['age']])

X_test[['workclass', 'occupation', 'native-country']] = imputer_cat.fit_transform(X_test[['workclass', 'occupation', 'native-country']])
X_test[['age']] = imputer_num.fit_transform(X_test[['age']])

3. 生成哑变量 #

# Create dummy vars
X_train = pd.get_dummies(X_train, columns=['workclass', 'education', 'marital-status', 
                                     'occupation', 'relationship', 'race', 'gender', 'native-country'], drop_first=True)
X_test = pd.get_dummies(X_test, columns=['workclass', 'education', 'marital-status', 
                                     'occupation', 'relationship', 'race', 'gender', 'native-country'], drop_first=True)

y_train = pd.get_dummies(y_train, columns='income', drop_first=True)
y_test = pd.get_dummies(y_test, columns='income', drop_first=True)

4. 展平目标变量 #

y_test = y_test.values.ravel()
y_train = y_train.values.ravel()

操作步骤：

• 将数据集拆分为训练集与测试集；
• 分别对训练集和测试集的缺失值进行填充；
• 为分类变量生成哑变量并处理目标变量。

执行后发现训练集与测试集列数不一致，部分列丢失。

原因分析 #

1. 哑变量生成逻辑独立：pd.get_dummies会根据当前数据集的类别生成列。如果测试集的分类变量存在训练集没有的类别，或者缺少训练集有的类别，就会导致两边的哑变量列数、列名不匹配。比如训练集的workclass包含Private、Self-emp-not-inc，而测试集只有Private，那测试集就会少一列对应的哑变量。

2. 测试集重新拟合填充器：对测试集的缺失值填充时，调用了imputer_cat.fit_transform(X_test)，这会重新基于测试集的数据拟合填充策略，而非沿用训练集拟合好的规则。这可能导致测试集填充后的类别分布与训练集不一致，进一步加剧哑变量列的差异。

3. 目标变量处理冗余：用pd.get_dummies处理二分类目标变量虽然可行，但不如LabelEncoder简洁，不过这不是列不一致的直接原因。

解决方案 #

核心原则：所有预处理逻辑必须基于训练集拟合，再应用到测试集，避免数据泄露和分布不一致。

修正步骤及代码 #

1. 正确处理缺失值

用训练集拟合的填充器转换测试集，禁止重新拟合：

from sklearn.impute import SimpleImputer
import pandas as pd
from sklearn.model_selection import train_test_split

random_state_value = 0

# 拆分数据集
X = data.drop(columns='income', axis=1)
y = data['income']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=random_state_value)

# 定义列类型
cat_cols = ['workclass', 'occupation', 'native-country']
num_cols = ['age']

# 拟合训练集填充器，转换训练集和测试集
imputer_cat = SimpleImputer(strategy='most_frequent')
imputer_cat.fit(X_train[cat_cols])
X_train[cat_cols] = imputer_cat.transform(X_train[cat_cols])
X_test[cat_cols] = imputer_cat.transform(X_test[cat_cols])

imputer_num = SimpleImputer(strategy='median')
imputer_num.fit(X_train[num_cols])
X_train[num_cols] = imputer_num.transform(X_train[num_cols])
X_test[num_cols] = imputer_num.transform(X_test[num_cols])

2. 对齐哑变量列

生成训练集哑变量后，用训练集的列对齐测试集，确保列数、列名完全一致：

# 生成训练集哑变量
X_train_dummies = pd.get_dummies(X_train, columns=['workclass', 'education', 'marital-status', 
                                     'occupation', 'relationship', 'race', 'gender', 'native-country'], drop_first=True)

# 生成测试集哑变量并对齐训练集列
X_test_dummies = pd.get_dummies(X_test, columns=['workclass', 'education', 'marital-status', 
                                     'occupation', 'relationship', 'race', 'gender', 'native-country'], drop_first=True)

# 对齐列，缺失的列填充0
X_test_dummies = X_test_dummies.reindex(columns=X_train_dummies.columns, fill_value=0)

# 更新变量
X_train = X_train_dummies
X_test = X_test_dummies

3. 优化目标变量处理（可选）

二分类目标变量用LabelEncoder更简洁高效：

from sklearn.preprocessing import LabelEncoder

le = LabelEncoder()
y_train = le.fit_transform(y_train)
y_test = le.transform(y_test)
# LabelEncoder输出为一维数组，无需再执行ravel操作

更规范的方案：使用Pipeline和ColumnTransformer #

用Sklearn的Pipeline整合所有预处理步骤，自动保证训练集和测试集的一致性：

from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import OneHotEncoder

# 定义数值型变量预处理管道
numeric_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='median'))
])

# 定义分类变量预处理管道
categorical_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='most_frequent')),
    ('onehot', OneHotEncoder(handle_unknown='ignore', drop='first'))
])

# 整合两类变量的预处理逻辑
preprocessor = ColumnTransformer(
    transformers=[
        ('num', numeric_transformer, num_cols),
        ('cat', categorical_transformer, cat_cols)
    ])

# 构建完整预处理管道（可后续添加模型）
clf = Pipeline(steps=[('preprocessor', preprocessor)])

# 拟合训练集，转换训练集和测试集
X_train_processed = clf.fit_transform(X_train)
X_test_processed = clf.transform(X_test)

# 处理目标变量
le = LabelEncoder()
y_train = le.fit_transform(y_train)
y_test = le.transform(y_test)

OneHotEncoder(handle_unknown='ignore')会自动忽略测试集中训练集未出现的类别，从根源避免列不一致问题。

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