scikit-learn Pipeline预测特征不匹配的Python优雅解决方案咨询

问题根因 #

scikit-learn的特征处理组件（尤其是ColumnTransformer、OneHotEncoder这类）在训练时会固定特征的顺序和数量，预测时要求输入的DataFrame列顺序和训练时完全一致，列缺失、多余、顺序不对都会触发特征维度/名称不匹配错误。

推荐的优雅解决方案 #

方案1：直接绑定训练特征名到流水线对象（最简实现） #

sklearn 1.0及以上版本的流水线在fit后会自动生成feature_names_in_属性存储训练时的输入列名，如果你的版本较低或者该属性缺失，可以在序列化前手动给流水线添加该属性，不需要额外存独立的列名文件：

build.py修改代码

test_pipeline.fit(X_train, y_train)
# 手动绑定训练特征名到流水线对象，会随pickle一起序列化
test_pipeline.feature_names_in_ = X_train.columns.tolist()

import pickle
# 注意：你原写法的dump参数顺序是反的，会直接序列化失败，第一个参数是待存储对象，第二个是文件句柄
# 推荐用with上下文管理器操作文件，避免资源泄漏
with open('final_pipeline.pkl','wb') as f:
    pickle.dump(test_pipeline, f)

app.py修改代码

import pickle
import pandas as pd

with open('final_pipeline.pkl', 'rb') as f:
    pipeline = pickle.load(f)

input_df = pd.DataFrame(data.dict(), index=[0])
# 用reindex一键完成3件事：按训练列顺序重排、缺失列补默认值、多余列自动丢弃
# 可根据业务需求调整fill_value，比如类别列填'unknown'、数值列填0
input_df = input_df.reindex(columns=pipeline.feature_names_in_, fill_value=None)

prediction = pipeline.predict(input_df)

方案2：自定义列对齐Transformer嵌入流水线（最符合sklearn规范） #

如果想要完全把列对齐逻辑封装到流水线里，app不需要做任何额外处理，可以自定义一个符合sklearn API规范的Transformer，放到流水线的最开头：

自定义列对齐组件

from sklearn.base import BaseEstimator, TransformerMixin

class FeatureAligner(BaseEstimator, TransformerMixin):
    def fit(self, X, y=None):
        # 训练阶段存储输入特征名
        self.feature_names_ = X.columns.tolist()
        return self
    
    def transform(self, X):
        # 预测阶段自动对齐列
        return X.reindex(columns=self.feature_names_, fill_value=None)

修改build.py的流水线定义

test_pipeline = Pipeline(steps=[
            # 把列对齐步骤放到流水线最开头
            ('feature_align', FeatureAligner()),
            ('preprocessor', preprocessor),
            ('std_scale', StandardScaler()),
            ('XGB_model', 
                xgb.XGBRegressor(
                    booster = 'gbtree', colsample_bylevel=0.75,colsample_bytree=0.75,
                    max_depth = 20,grow_policy = 'depthwise',learning_rate = 0.1
                 )
             )
        ])

这种方式序列化后的流水线可以直接调用，app里不需要加任何额外的列处理逻辑：

pipeline = pickle.load(open('final_pipeline.pkl', 'rb'))
input_df = pd.DataFrame(data.dict(), index=[0])
prediction = pipeline.predict(input_df)

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