sklearn Pipeline中仅缓存UMAP单步骤 解决自定义Transformer序列化报错

方案1：封装带独立缓存的UMAP自定义Transformer #

直接继承UMAP类，在fit逻辑里加入基于joblib的本地缓存判断，不需要依赖Pipeline的全局memory机制，也不会受上游不可序列化Transformer的影响，同时可以保留Pipeline的完整调用形式，适配GridSearchCV等sklearn工具链。

代码实现如下：

import joblib
import hashlib
import numpy as np
import os
from umap import UMAP

class CachedUMAP(UMAP):
    def __init__(self, cache_dir='./umap_cache', **kwargs):
        super().__init__(**kwargs)
        self.cache_dir = cache_dir
        os.makedirs(cache_dir, exist_ok=True)
    
    def fit(self, X, y=None):
        # 计算输入数据哈希，避免不同输入复用缓存
        data_hash = hashlib.sha256(np.ascontiguousarray(X)).hexdigest()[:16]
        # 计算超参数哈希，超参数变更时自动更新缓存
        params_hash = hashlib.sha256(str(self.get_params()).encode()).hexdigest()[:16]
        cache_path = f"{self.cache_dir}/umap_{data_hash}_{params_hash}.pkl"
        
        if os.path.exists(cache_path):
            # 存在缓存则直接加载模型参数
            cached_obj = joblib.load(cache_path)
            self.__dict__.update(cached_obj.__dict__)
            return self
        
        # 无缓存则正常训练后保存
        super().fit(X, y)
        joblib.dump(self, cache_path)
        return self

替换后的Pipeline调用代码：

from sklearn.preprocessing import FunctionTransformer
from sklearn.pipeline import Pipeline
from hdbscan import HDBSCAN
import seaborn as sns

iris = sns.load_dataset("iris")
X = iris.drop(columns='species')
y = iris.species

@FunctionTransformer
def transform_something(iris):
    iris = iris.copy()
    iris['sepal_sum'] = iris.sepal_length + iris.sepal_width
    return iris

# 不需要配置全局memory参数
pipe = Pipeline([
                 ('transformer', transform_something),
                 ('umap', CachedUMAP()), # 替换为带缓存的UMAP类
                 ('hdb', HDBSCAN()),
                ]
            )

pipe.fit_predict(X)

方案2：拆分Pipeline手动控制缓存 #

如果不想自定义类，也可以拆分Pipeline为三个独立环节，手动控制UMAP的缓存逻辑：

1. 运行上游自定义Transformer得到预处理后的数据
2. 单独判断UMAP缓存是否存在，存在则直接加载，不存在则训练后保存
3. 运行下游HDBSCAN步骤

代码实现如下：

from sklearn.preprocessing import FunctionTransformer
from umap import UMAP
from hdbscan import HDBSCAN
import seaborn as sns
import joblib
import os
import hashlib
import numpy as np

# 数据加载
iris = sns.load_dataset("iris")
X = iris.drop(columns='species')
y = iris.species

# 1. 执行上游自定义转换
@FunctionTransformer
def transform_something(iris):
    iris = iris.copy()
    iris['sepal_sum'] = iris.sepal_length + iris.sepal_width
    return iris
X_transformed = transform_something.fit_transform(X)

# 2. 单独处理UMAP缓存逻辑
cache_dir = "./umap_cache"
cache_path = f"{cache_dir}/umap_model.pkl"
os.makedirs(cache_dir, exist_ok=True)

if os.path.exists(cache_path):
    umap = joblib.load(cache_path)
    X_umap = umap.transform(X_transformed)
else:
    umap = UMAP()
    X_umap = umap.fit_transform(X_transformed)
    joblib.dump(umap, cache_path)

# 3. 执行下游聚类
hdb = HDBSCAN()
result = hdb.fit_predict(X_umap)

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