如何将两个具有 IF 逻辑条件的子 scikit-learn 模型包装成一个集成模型并将其保存到 pickle 文件中？

Question

我使用 IF 条件训练了两个 scikit-learn 模型（根据 X1 特征定义的标准生成 2 个训练集）。如何将这个 IF 条件与这两个经过训练的模型包装成一个整体模型并将其保存到 pickle 模型文件中？谢谢！

Answer 1

如果我理解正确的话，您已经训练了两个单独的模型，两个模型都针对训练数据的子集（例如儿童/成人）进行了训练，并且每个模型都是整个数据集的一部分的专家？现在你的问题是，是否有可能有一个类似“路由实例”的东西，在你进行预测时将样本传递给两个模型中的任何一个？

如果是这样，你可以试试这个：

import pickle
from sklearn.base import BaseEstimator

class EnsembleModel(BaseEstimator):
    def __init__(self, model1, model2, threshold):
        self.model1 = model1
        self.model2 = model2
        self.threshold = threshold

    def predict(self, X):
        if X['X1'] < self.threshold:
            return self.model1.predict(X)
        else:
            return self.model2.predict(X)

# Usage:
ensemble = EnsembleModel(model1, model2, threshold=18)

# Saving:
with open('ensemble_model.pkl', 'wb') as file:
    pickle.dump(ensemble, file)

# Loading:
with open('ensemble_model.pkl', 'rb') as file:
    loaded_ensemble = pickle.load(file)