如何在scikit-learn的Pipeline流水线中集成使用t-SNE算法

问题根因 #

Sklearn原生TSNE类仅提供fit_transform方法，没有实现独立的transform接口，不符合Pipeline要求中间步骤必须同时具备fit、transform能力的规则，因此必须自定义符合接口规范的转换器。
此外需要注意：t-SNE算法本身是基于输入样本的全局邻接关系生成低维嵌入，不存在训练完成后可直接套用到新样本的映射函数，因此自定义转换器的transform能力是通过近似方法实现的。

自定义转换器实现 #

你需要继承Sklearn的BaseEstimator和TransformerMixin基类，通过K近邻近似实现新样本的低维映射，完整实现代码如下：

from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.manifold import TSNE
from sklearn.neighbors import NearestNeighbors
import numpy as np

class TSNETransformer(BaseEstimator, TransformerMixin):
    def __init__(self, n_components=2, perplexity=30, random_state=42, n_neighbors=5, **tsne_kwargs):
        # 接收TSNE的所有参数
        self.n_components = n_components
        self.perplexity = perplexity
        self.random_state = random_state
        self.tsne_kwargs = tsne_kwargs
        # transform阶段用的近邻数量
        self.n_neighbors = n_neighbors

    def fit(self, X, y=None):
        # 拟合TSNE得到训练集的低维嵌入
        self.tsne = TSNE(
            n_components=self.n_components,
            perplexity=self.perplexity,
            random_state=self.random_state,
            **self.tsne_kwargs
        )
        self.train_embedding_ = self.tsne.fit_transform(X)
        # 存储训练集原始特征，用于后续近邻查询
        self.X_train_ = X
        # 拟合近邻查询模型
        self.nn_ = NearestNeighbors(n_neighbors=self.n_neighbors)
        self.nn_.fit(X)
        return self

    def transform(self, X):
        # 查找每个新样本在训练集中的k个近邻
        distances, indices = self.nn_.kneighbors(X)
        # 取近邻的低维嵌入的均值作为新样本的嵌入结果
        return np.mean(self.train_embedding_[indices], axis=1)

结合Pipeline与GridSearchCV的使用示例 #

你可以直接把自定义转换器放到Pipeline的步骤中，和其他组件串联使用：

from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import GridSearchCV

# 构建Pipeline
pipe = Pipeline([
    ('scaler', StandardScaler()),
    ('tsne', TSNETransformer(random_state=42)),
    ('rf', RandomForestClassifier(random_state=42))
])

# 定义参数搜索网格
param_grid = {
    'tsne__n_components': [2, 3],
    'tsne__perplexity': [20, 30, 50],
    'rf__n_estimators': [100, 200],
    'rf__max_depth': [None, 10, 20]
}

# 初始化网格搜索
grid = GridSearchCV(pipe, param_grid, cv=5, scoring='accuracy', n_jobs=-1)
# 训练，输入原始特征即可
grid.fit(X, y)

# 输出最优参数和得分
print("最优参数：", grid.best_params_)
print("最优交叉验证得分：", grid.best_score_)

注意事项 #

• t-SNE计算复杂度极高，数据量大于1万条时建议先加一层PCA降维到30~50维后再输入TSNE，可大幅降低运行耗时
• 基于K近邻的近似变换存在一定精度损失，如果你的场景对降维精度要求极高，建议先固定t-SNE超参数生成降维结果，再单独对下游分类器做参数调优
• 禁止对全量数据集做TSNE降维后再拆分训练测试集，该操作会引入严重的数据泄露，导致模型评估结果虚高、完全不可信

内容的提问来源于stack exchange，提问作者Mr. Hankey