Python sklearn模型是什么类型的？

我正在编写一些代码，根据一些数据评估不同的sklearn模型。我正在使用类型提示，这既是为了我自己的教育，也是为了帮助其他最终不得不阅读我的代码的人

我的问题是如何指定sklearn预测器的类型（例如

LinearRegression（）

）

例如：

def model_tester(model : Predictor,
                 parameter: int
                 ) -> np.ndarray:
     """An example function with type hints."""

     # do stuff to model 

     return values

我看到可以创建新类型，也可以使用

TypeVar

执行以下操作：

Predictor = TypeVar('Predictor') 

但如果sklearn模型已经有了传统的类型，我就不想使用这个

检查LinearRegression（）的类型会产生：

---

这显然是有用的，但前提是我对线性回归模型感兴趣

我认为所有模型继承的最通用类是

sklearn.base.BaseEstimator

如果您想更具体一些，可以使用

sklearn.base.ClassifierMixin

或

sklearn.base.RegressionrMixin

因此，我会：

from sklearn.base import RegressorMixin


def model_tester(model: RegressorMixin, parameter: int) -> np.ndarray:
     """An example function with type hints."""

     # do stuff to model 

     return values

---

我不是类型检查方面的专家，如果这不正确，请纠正我。

一个很好的解决方法是创建您自己的自定义类型提示类（使用Union），其中包括您常用的所有模型。它需要更多的努力，但可以让你更具体，并与PyCharm合作

ModelRegressor = Union[LinearRegression, DecisionTreeRegressor, RandomForestRegressor, SVR]

def foo(model: ModelRegressor):
    do_something

在Python3.8上（或更早版本的使用），您可以使用

键入.Protocol

。使用协议，您可以使用一个名为的概念来准确定义类型的预期结构：

from typing import Protocol
# from typing_extensions import Protocol  # for Python <3.8

class ScikitModel(Protocol):
    def fit(self, X, y, sample_weight=None): ...
    def predict(self, X): ...
    def score(self, X, y, sample_weight=None): ...
    def set_params(self, **params): ...

---

谢谢你的回答。我尝试了BaseEstimator和ClassifierMixin。但是当我调用self.estimator.fit时，我的IDE（Pycharm）抱怨它找不到“fit”属性。这是正确的。这些类不实现fit。对于每个估计器（例如LogisticReturnal），它是单独实施的。有人知道应该采用scikit学习估计器的参数的正确类型提示是什么吗？

from typing import Protocol
# from typing_extensions import Protocol  # for Python <3.8

class ScikitModel(Protocol):
    def fit(self, X, y, sample_weight=None): ...
    def predict(self, X): ...
    def score(self, X, y, sample_weight=None): ...
    def set_params(self, **params): ...

def do_stuff(model: ScikitModel) -> Any:
    model.fit(train_data, train_labels)  # this type checks 
    score = model.score(test_data, test_labels)  # this type checks
    ...