使用set_output API配置Transformer输出Pandas DataFrame

在数据分析和机器学习中，Pandas DataFrame是一种非常强大的数据结构，它允许以表格形式组织和操作数据。Scikit-learn是一个流行的机器学习库，它提供了许多用于数据预处理、模型训练和评估的工具。为了更好地集成Pandas DataFrame与scikit-learn的工作流，可以利用set_outputAPI来配置Transformer，使其输出Pandas DataFrame。本文将详细介绍如何使用这一API。

加载数据集

首先，需要加载一个数据集来演示set_outputAPI的使用。这里使用iris数据集，它是一个经典的多类分类数据集，包含了150个样本，每个样本有4个特征。

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

# 加载iris数据集，并将其转换为DataFrame
X, y = load_iris(as_frame=True, return_X_y=True)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y, random_state=0)

# 显示训练集的前几行数据
print(X_train.head())

配置Transformer输出Pandas DataFrame

要配置一个estimator（如StandardScaler）以返回DataFrame，可以调用set_output方法。此外，也可以通过设置全局配置transform_output="pandas"来实现。

from sklearn.preprocessing import StandardScaler

# 创建StandardScaler实例
scaler = StandardScaler()

# 配置StandardScaler输出DataFrame
scaler.set_output(transform="pandas")

# 训练模型并转换测试集数据
scaler.fit(X_train)
X_test_scaled = scaler.transform(X_test)

# 显示转换后的数据
print(X_test_scaled.head())

在Pipeline中使用set_output

在scikit-learn中，Pipeline是一种将多个处理步骤串联起来的工具。通过在Pipeline中使用set_output，可以确保所有步骤都输出DataFrame。

from sklearn.feature_selection import SelectPercentile
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline

# 创建一个Pipeline，包含标准化、特征选择和逻辑回归
clf = make_pipeline(StandardScaler(), SelectPercentile(percentile=75), LogisticRegression())

# 配置Pipeline输出DataFrame
clf.set_output(transform="pandas")

# 训练模型
clf.fit(X_train, y_train)

全局配置set_output

除了在单个estimator上使用set_output外，还可以全局配置transform_output为"pandas"，这样所有estimator都会输出DataFrame。

from sklearn import set_config
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import OneHotEncoder, StandardScaler

# 全局配置transform_output为"pandas"
set_config(transform_output="pandas")

# 创建一个ColumnTransformer，包含数值和分类特征的处理
num_pipe = make_pipeline(SimpleImputer(), StandardScaler())
num_cols = ["age", "fare"]
ct = ColumnTransformer((
    ("numerical", num_pipe, num_cols),
    ("categorical", OneHotEncoder(sparse_output=False, drop="if_binary", handle_unknown="ignore"), ["embarked", "sex", "pclass"]),
), verbose_feature_names_out=False)

# 创建一个Pipeline，包含ColumnTransformer和逻辑回归
clf = make_pipeline(ct, SelectPercentile(percentile=50), LogisticRegression())

# 训练模型并评估测试集
clf.fit(X_train, y_train)
print(clf.score(X_test, y_test))

使用config_context

config_context是一个上下文管理器，它允许在特定的代码块中临时更改配置。这对于在不同情况下使用不同的输出格式非常有用。

from sklearn import config_context

# 创建一个StandardScaler实例
scaler = StandardScaler()

# 在config_context中配置transform_output为"pandas"
with config_context(transform_output="pandas"):
    # 在这个上下文中，transform的输出将是一个Pandas DataFrame
    X_test_scaled = scaler.transform(X_test[num_cols])
    print(X_test_scaled.head())

# 在上下文外部，输出将是一个NumPy数组
X_test_scaled = scaler.transform(X_test[num_cols])
print(X_test_scaled[:5])