要将前k个结果添加到sklearn管道？

Question

我想将我的模型放在云中，但是它需要一个 joblib 文件。 我有一个模型，该模型可以预测类（多类分类）。 我目前使用np.argsort并获得最高的k结果。

如何在sklearn管道中添加另一步骤来获取模型的输出，以便单个管道将 top k 结果返回给我？

由于AFAIK， joblib 只能转储经过训练的模型的单个管道。

Answer 1

您可以继承Pipeline，并添加一种使用概率来获取top k预测的方法。

from sklearn.pipeline import Pipeline
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.datasets import make_classification

X, y = make_classification(n_features=20, n_informative=15, n_classes=5)

class MyPipeline(Pipeline):
    def get_top_k(self, X, k=3):
        probs = self.predict_proba(X)
        return np.fliplr(np.argsort(probs, axis=1)[:, -k:])

X_train, X_test, y_train, y_test = train_test_split(X, y,
                                                    random_state=42)
pipe = MyPipeline([('scalling', StandardScaler()),
                   ('logistic', LogisticRegression())])
pipe.fit(X_train, y_train)
pipe.get_top_k(X_test[:2, :])

现在使用pipe转储joblib对象。