如何在Scikit-learn Pipeline中跨步骤传递拟合参数？

我正在编写自定义Scikit-learn Pipeline：

Step 1 代码：

class Step1(BaseEstimator, TransformerMixin):
    def __init__(self, input1: str = "Input1") -> None:
        self.input1 = input1

    def fit(self, X: pd.DataFrame, y: pd.DataFrame = None):
        self.some_param = X["my_col"].mean()
        return self
    
    def transform(self, X):
        return X

Step 2 代码：

class Step2(BaseEstimator, TransformerMixin):
    def __init__(self, input2: str = "Input2") -> None:
        self.input2 = input2

    def fit(self, X: pd.DataFrame, y: pd.DataFrame = None):
        return self
    
    def transform(self, X):
        # 想要在此处使用Step1中拟合得到的变量: 'some_param'
        X["my_other_col"] = X["my_other_col"] - (some_param)
        return X

Pipeline定义：

pipeline = Pipeline([
('step1', Step1()),
('step2', Step2())
])

需求：在Step1中拟合得到的some_param，需要在Step2的transform方法中使用，尝试过set_param但效果不符合预期，求可行实现方法。

可行实现方法 #

方法一：拟合后手动传递参数 #

修改Step2，添加some_param作为初始化参数，在Pipeline拟合完成后，从Step1实例中取出参数赋值给Step2：

class Step2(BaseEstimator, TransformerMixin):
    def __init__(self, input2: str = "Input2", some_param: float = None) -> None:
        self.input2 = input2
        self.some_param = some_param  # 添加参数存储属性

    def fit(self, X: pd.DataFrame, y: pd.DataFrame = None):
        return self
    
    def transform(self, X):
        if self.some_param is None:
            raise ValueError("请先拟合Pipeline并设置some_param参数")
        X["my_other_col"] = X["my_other_col"] - self.some_param
        return X

使用方式：

pipeline = Pipeline([
    ('step1', Step1()),
    ('step2', Step2())
])

# 拟合Pipeline
pipeline.fit(X_train)

# 将Step1的some_param传递给Step2
pipeline.named_steps['step2'].some_param = pipeline.named_steps['step1'].some_param

# 执行转换
result = pipeline.transform(X_test)

方法二：通过DataFrame传递参数 #

让Step1在transform时将some_param作为临时列加入DataFrame，Step2读取后再移除该列，无需手动传递：

修改Step1：

class Step1(BaseEstimator, TransformerMixin):
    def __init__(self, input1: str = "Input1") -> None:
        self.input1 = input1

    def fit(self, X: pd.DataFrame, y: pd.DataFrame = None):
        self.some_param = X["my_col"].mean()
        return self
    
    def transform(self, X):
        X_copy = X.copy()  # 避免修改原始数据
        X_copy['_temp_some_param'] = self.some_param  # 添加临时列
        return X_copy

修改Step2：

class Step2(BaseEstimator, TransformerMixin):
    def __init__(self, input2: str = "Input2") -> None:
        self.input2 = input2

    def fit(self, X: pd.DataFrame, y: pd.DataFrame = None):
        return self
    
    def transform(self, X):
        X_copy = X.copy()
        # 从临时列读取参数（所有行值相同，取第一行即可）
        some_param = X_copy['_temp_some_param'].iloc[0]
        X_copy["my_other_col"] = X_copy["my_other_col"] - some_param
        X_copy.drop('_temp_some_param', axis=1, inplace=True)  # 删除临时列
        return X_copy

使用方式：

pipeline = Pipeline([
    ('step1', Step1()),
    ('step2', Step2())
])

pipeline.fit(X_train)
result = pipeline.transform(X_test)

方法三：自定义复合Transformer #

将Step1和Step2封装到一个类中，内部直接传递参数，简化Pipeline结构：

class CombinedTransformer(BaseEstimator, TransformerMixin):
    def __init__(self, input1: str = "Input1", input2: str = "Input2") -> None:
        self.input1 = input1
        self.input2 = input2
        self.step1 = Step1(input1)
        self.step2 = Step2(input2)

    def fit(self, X: pd.DataFrame, y: pd.DataFrame = None):
        self.step1.fit(X, y)
        # 内部直接传递参数
        self.step2.some_param = self.step1.some_param
        self.step2.fit(X, y)
        return self
    
    def transform(self, X):
        X = self.step1.transform(X)
        X = self.step2.transform(X)
        return X

使用方式：

pipeline = Pipeline([
    ('combined', CombinedTransformer())
])

pipeline.fit(X_train)
result = pipeline.transform(X_test)

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