如何根据DataFrame列值应用不同预测模型，或合并模型避免索引问题？

问题背景 #

现有用于血压预测的数据集（示例如下），已按性别拆分并分别训练了不同模型，但对新数据预测时，拆分后应用模型再合并会破坏原索引，需要解决这个问题。

示例数据集 #

data
    #  weight height age  Hours_active   sex   Blood_Pressure
    #1     74   1.68  22     5         Male    0.8
    #2     85   1.83  25    7          Male    0.95
    #3     58   1.58  21    2          Woman   0.97
    #4     80   1.72  20    4          Woman   0.23
    #5     70   1.72  24    1          Woman   0.40

已完成的数据集拆分与模型训练 #

拆分数据集

data_males = data[data['sex'] == "Male"].drop(['sex'], axis=1)
data_females = data[data['sex'] == "Woman"].drop(['sex'], axis=1)

训练模型

from sklearn.linear_model import LinearRegression, SGDRegressor

# 男性模型训练
X_train_m= data_males.drop(['Blood_Pressure'], axis=1)
y_train_m =  data_males['Blood_Pressure']
lin= LinearRegression()
model_males= lin.fit(X_train_m, y_train_m)

# 女性模型训练
X_train_f= data_females.drop(['Blood_Pressure'],axis=1)
y_train_f =  data_females['Blood_Pressure']
SGD = SGDRegressor()
model_females= SGD.fit(X_train_f, y_train_f)

解决方案 #

方案一：按性别直接应用对应模型，保留原索引 #

用apply逐行处理新数据，根据每行的sex值调用对应模型预测，全程保留原索引结构，无需拆分合并。

示例代码：

def predict_blood_pressure(row):
    # 提取特征（需与训练时的特征顺序一致）
    features = row[['weight', 'height', 'age', 'Hours_active']].values.reshape(1, -1)
    if row['sex'] == 'Male':
        return model_males.predict(features)[0]
    elif row['sex'] == 'Woman':
        return model_females.predict(features)[0]
    else:
        # 处理未知性别，可返回NaN或自定义值
        return None

# 对新数据集df_new执行预测
df_new['predicted_BP'] = df_new.apply(predict_blood_pressure, axis=1)

方案二：合并为一个考虑性别的模型 #

如果想把两个子模型整合为一个统一模型，有两种实现方式：

方式1：将性别作为特征训练统一模型

把性别做独热编码后加入特征集，训练一个全局模型，让模型自动学习性别与其他特征的交互关系：

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

# 定义特征预处理流程：对性别列做独热编码，其他特征直接保留
preprocessor = ColumnTransformer(
    transformers=[
        ('cat', OneHotEncoder(sparse_output=False), ['sex'])
    ],
    remainder='passthrough'
)

# 构建完整训练流水线
full_pipeline = Pipeline([
    ('preprocess', preprocessor),
    ('model', LinearRegression())  # 也可替换为SGDRegressor等其他模型
])

# 用完整数据集训练
X_full = data.drop(['Blood_Pressure'], axis=1)
y_full = data['Blood_Pressure']
full_pipeline.fit(X_full, y_full)

# 预测新数据时直接调用
df_new['predicted_BP'] = full_pipeline.predict(df_new)

方式2：自定义封装双模型的类

如果要保留原有的分性别训练逻辑，可自定义一个符合scikit-learn接口的模型类，对外提供统一的fit和predict方法：

from sklearn.base import BaseEstimator, RegressorMixin

class GenderSplitModel(BaseEstimator, RegressorMixin):
    def __init__(self, male_model, female_model):
        self.male_model = male_model
        self.female_model = female_model
    
    def fit(self, X, y):
        # 拆分数据分别训练两个子模型
        X_male = X[X['sex'] == 'Male'].drop(['sex'], axis=1)
        y_male = y[X['sex'] == 'Male']
        self.male_model.fit(X_male, y_male)
        
        X_female = X[X['sex'] == 'Woman'].drop(['sex'], axis=1)
        y_female = y[X['sex'] == 'Woman']
        self.female_model.fit(X_female, y_female)
        return self
    
    def predict(self, X):
        # 复制原数据避免修改
        X_copy = X.copy()
        # 分别对男性、女性数据做预测
        mask_male = X_copy['sex'] == 'Male'
        X_male = X_copy[mask_male].drop(['sex'], axis=1)
        X_copy.loc[mask_male, 'pred'] = self.male_model.predict(X_male)
        
        mask_female = X_copy['sex'] == 'Woman'
        X_female = X_copy[mask_female].drop(['sex'], axis=1)
        X_copy.loc[mask_female, 'pred'] = self.female_model.predict(X_female)
        
        # 返回预测结果，保持原数据的顺序与索引
        return X_copy['pred'].values

# 初始化并训练自定义模型
custom_model = GenderSplitModel(LinearRegression(), SGDRegressor())
custom_model.fit(X_full, y_full)

# 预测新数据
df_new['predicted_BP'] = custom_model.predict(df_new)

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