使用Pipeline.predict()传入列表触发KeyError的解决方法

问题根源 #

你用的category_encoders.OneHotEncoder是基于DataFrame列名工作的——训练时它会记住每个分类特征对应的列名（比如报错里的sex），而直接传入列表/数组时，数据没有列名标识，编码器找不到对应的特征列，因此触发KeyError。

解决方案 #

不管是本地测试还是FastAPI部署，核心思路都是给输入数据补上训练时的特征列名，转成DataFrame再传入Pipeline。

方法1：手动将列表转为带列名的DataFrame（本地测试用）

先准备好训练集的特征列名（必须和X_train的列名完全一致，顺序也要对应），再把输入列表包装成DataFrame：

import pandas as pd

# 替换成你真实的25个特征列名，顺序要和训练时X_train的列顺序一致
feature_names = ['sex', 'age', 'education', 'workclass', 'marital_status', 
                 'fnlwgt', 'education_num', 'occupation', 'relationship', 
                 'race', 'capital_gain', 'capital_loss', 'hours_per_week',
                 'native_country', ...]  # 补全剩余特征名

# 输入列表转DataFrame
input_list = ['M', 15, 'U', 'LE3', 'T', 4, 3, 'teacher', 'services', 1, 3, 0,
              'no', 'yes', 'no', 'yes', 'yes', 'yes', 'yes', 5, 4, 4, 2, 15, 16]
input_df = pd.DataFrame([input_list], columns=feature_names)

# 正常执行预测
pipe.predict(input_df)

方法2：FastAPI部署用Pydantic模型规范输入（生产环境推荐）

用Pydantic定义输入数据结构，自动校验格式并映射列名，避免手动维护列名顺序：

from fastapi import FastAPI
from pydantic import BaseModel
import pandas as pd

app = FastAPI()
# 假设你的pipe已经加载完成（可以用joblib/pickle提前保存模型，启动时加载）

# 定义输入模型，字段名必须和训练时的特征列名完全一致
class PredictionInput(BaseModel):
    sex: str
    age: int
    education: str
    workclass: str
    marital_status: str
    fnlwgt: int
    education_num: int
    occupation: str
    relationship: str
    race: str
    capital_gain: int
    capital_loss: int
    hours_per_week: int
    native_country: str
    # 补全剩余11个特征的字段定义

@app.post("/predict")
def make_prediction(input_data: PredictionInput):
    # 将Pydantic模型转为DataFrame
    input_df = pd.DataFrame([input_data.dict()])
    # 执行预测
    pred_result = pipe.predict(input_df)[0]
    return {"prediction": float(pred_result)}

调用API时，只需传入JSON格式的请求体：

{
    "sex": "M",
    "age": 15,
    "education": "U",
    "workclass": "LE3",
    // 其他特征字段...
}

关键注意事项 #

• 特征列名必须和训练集X_train的列名完全匹配，包括大小写、拼写，不能有任何差异
• 如果用列表转DataFrame，特征顺序必须和训练集一致；用Pydantic模型则无需担心顺序，因为是按字段名映射
• 永远不要直接给依赖列名的编码器传入无列名的数组/列表，这会破坏训练时的特征映射逻辑

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