Scikit-learn FeatureUnion结合PyTorch：并行任务共享大模型避免复制的方法

这个场景太常见了——大模型在多进程里被重复加载，显存/内存直接炸掉对吧？核心问题是Scikit-learn的n_jobs默认用多进程fork，每个进程会完整复制一份转换器对象（包括里面的PyTorch模型），完全没有共享可言。下面给你几个实用的解决方案，按可行性排序：

1. 合并两个转换器为单一Transformer（最简单） #

如果两个转换器只是对同一模型的不同输出做特征提取（比如取不同层的输出），直接把逻辑合并成一个Transformer是最省心的：

from sklearn.base import BaseEstimator, TransformerMixin
import torch
import numpy as np

class CombinedTransformer(BaseEstimator, TransformerMixin):
    def __init__(self, net):
        self.net = net
        # 确保模型处于评估模式，关闭dropout等训练层
        self.net.eval()

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        with torch.no_grad():
            # 一次前向传播拿到所有需要的特征
            outputs = self.net(X)
            feat1 = outputs["feat1"].cpu().numpy()
            feat2 = outputs["feat2"].cpu().numpy()
        # 拼接特征，和FeatureUnion效果完全一致
        return np.hstack([feat1, feat2])

# 使用时直接替换FeatureUnion
combined_transformer = CombinedTransformer(net)
# 后续可以直接用在Pipeline里，不用n_jobs也能高效运行（一次前向传播搞定两种特征）

这样不管用不用多进程，都只有一个模型实例，完全避免复制问题。

2. 用线程代替多进程（GPU场景首选） #

如果必须保留两个转换器，且模型跑在GPU上，Scikit-learn 0.23+版本支持指定backend="threading"来替换默认的多进程：

from sklearn.pipeline import FeatureUnion

union = FeatureUnion(
    [('transformer_1', transformer_1), ('transformer_2', transformer_2)],
    n_jobs=2,
    backend="threading"
)

原因是：GPU推理操作会自动释放Python GIL，线程之间可以共享同一个GPU模型实例，不会复制权重。而且线程的开销比进程小得多，适合计算密集型的GPU任务。

⚠️ 注意：如果模型跑在CPU上，线程模式可能因为GIL导致速度变慢，这时候看下面的共享内存方案。

3. 用PyTorch共享内存存储模型权重（CPU场景适用） #

对于CPU上的大模型，可以利用PyTorch的torch.multiprocessing把模型权重放到共享内存，让所有子进程共享同一份权重，而不是复制。

第一步：主进程初始化模型并把权重转成共享张量 #

import torch.multiprocessing as mp

# 主进程初始化模型
net = SomePytorchModule()
net.eval()

# 把模型的state_dict转成共享张量
shared_state_dict = {}
for k, v in net.state_dict().items():
    # 把张量移到共享内存，子进程可以直接访问
    shared_state_dict[k] = v.share_memory_()

第二步：修改转换器，从共享内存加载权重 #

class SharedTransformer1(BaseEstimator, TransformerMixin):
    def __init__(self, shared_state_dict):
        self.shared_state_dict = shared_state_dict
        # 延迟初始化模型，避免主进程提前复制
        self.net = None

    def fit(self, X, y=None):
        # 在子进程里初始化模型并加载共享权重
        if self.net is None:
            self.net = SomePytorchModule()
            self.net.load_state_dict(self.shared_state_dict)
            self.net.eval()
        return self

    def transform(self, X):
        with torch.no_grad():
            feat = self.net(X).cpu().numpy()
        return feat

# 同理定义SharedTransformer2...

# 创建转换器实例，传入共享权重字典
transformer_1 = SharedTransformer1(shared_state_dict)
transformer_2 = SharedTransformer2(shared_state_dict)

# 然后正常使用FeatureUnion + n_jobs=2
union = FeatureUnion([('t1', transformer_1), ('t2', transformer_2)], n_jobs=2)

这样每个子进程会自己初始化模型结构，但权重是从共享内存加载的，不会复制大张量，内存占用会大幅降低。

4. 自定义多进程并行逻辑（最灵活但复杂） #

如果上面的方法都不满足，可以用torch.multiprocessing.Pool自己实现并行逻辑，完全掌控流程：

def process_transformer(transformer_cls, shared_state_dict, X):
    transformer = transformer_cls(shared_state_dict)
    transformer.fit(X)
    return transformer.transform(X)

# 主进程准备数据和共享权重
X = ... # 你的输入数据
shared_state_dict = ... # 同方案3的共享权重字典

# 创建进程池
with mp.Pool(2) as pool:
    # 并行执行两个转换器
    feat1, feat2 = pool.starmap(
        process_transformer,
        [(SharedTransformer1, shared_state_dict, X), (SharedTransformer2, shared_state_dict, X)]
    )

# 拼接特征得到最终结果
final_feat = np.hstack([feat1, feat2])

这种方式适合复杂场景，但需要自己处理数据拆分、结果拼接等细节。

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