如何在sklearn pipeline中添加列运算步骤实现两列相乘并删除原列

这个需求可以通过自定义sklearn转换器实现，只需要继承sklearn提供的BaseEstimator和TransformerMixin基类，即可无缝接入pipeline流程。

步骤1：自定义列相乘转换器 #

import pandas as pd
import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin

class ColumnMultiplier(BaseEstimator, TransformerMixin):
    def __init__(self, col1, col2, drop_original=True):
        """
        初始化参数
        :param col1: 要相乘的第一列，DataFrame传入列名，numpy数组传入索引编号
        :param col2: 要相乘的第二列，格式同上
        :param drop_original: 是否删除原始的两列，默认是
        """
        self.col1 = col1
        self.col2 = col2
        self.drop_original = drop_original
    
    def fit(self, X, y=None):
        # 该转换无训练参数，直接返回自身
        return self
    
    def transform(self, X):
        X_copy = X.copy()
        # 兼容DataFrame和numpy数组输入
        if isinstance(X_copy, pd.DataFrame):
            X_copy['product_col'] = X_copy[self.col1] * X_copy[self.col2]
            if self.drop_original:
                X_copy = X_copy.drop(columns=[self.col1, self.col2])
        else:
            product = X_copy[:, self.col1] * X_copy[:, self.col2]
            product = product.reshape(-1, 1)
            if self.drop_original:
                # 保留未参与运算的列，拼接乘积列
                keep_idx = [i for i in range(X_copy.shape[1]) if i not in [self.col1, self.col2]]
                X_copy = np.hstack([X_copy[:, keep_idx], product])
            else:
                X_copy = np.hstack([X_copy, product])
        return X_copy

步骤2：整合到Pipeline使用 #

from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression
from sklearn.pipeline import Pipeline
from sklearn.model_selection import train_test_split

# 测试数据准备
df = pd.DataFrame({
    'A': np.random.rand(100),
    'B': np.random.rand(100),
    'C': np.random.rand(100),
    'Y': np.random.rand(100)
})
X = df[['A', 'B', 'C']]
Y = df['Y']
x_train, x_test, y_train, y_test = train_test_split(X, Y, train_size=0.8, random_state=42)

# 构建pipeline
pipe = Pipeline([
    # 如果输入是DataFrame，也可以传列名：ColumnMultiplier('B', 'C')
    ('product', ColumnMultiplier(col1=1, col2=2)), 
    ('scaler', StandardScaler()),
    ('linear', LinearRegression())
])

# 训练
pipe.fit(x_train, y_train)

# 直接传入原始特征预测，不需要手动计算乘积
print(pipe.predict(np.array([[1,3,4]])))

轻量替代方案：使用FunctionTransformer #

如果不需要高度自定义参数，也可以用sklearn自带的FunctionTransformer快速实现：

from sklearn.preprocessing import FunctionTransformer

def multiply_bc(X):
    # 适配输入特征顺序为A、B、C的场景
    return np.hstack([X[:, [0]], (X[:, 1] * X[:, 2]).reshape(-1, 1)])

pipe = Pipeline([
    ('product', FunctionTransformer(multiply_bc)),
    ('scaler', StandardScaler()),
    ('linear', LinearRegression())
])

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