Writing custom estimators

Custom logic that breaks the contract fails subtly in cross-validation and grid search.

To plug your own code into pipelines and searches, write a class that follows the estimator contract:

Inherit BaseEstimator (gives get_params, set_params, a repr and HTML display) and a mixin: TransformerMixin (adds fit_transform), ClassifierMixin or RegressorMixin (add score and tags). Mixins go before BaseEstimator in the bases.
__init__ only stores its arguments, unchanged, as attributes with the same names. No validation, no defaults computed from data. clone and set_params rely on this.
fit(X, y=None) validates input, learns state into attributes ending in _, and returns self.
transform / predict check the estimator is fitted (check_is_fitted(self)) and use only the learned state.

Validate with sklearn.utils.validation.validate_data(self, X, reset=True) in fit and reset=False afterwards: it converts input, records n_features_in_ and feature_names_in_, and checks consistency. Adding OneToOneFeatureMixin gives get_feature_names_out for column-preserving transformers.

sklearn.utils.estimator_checks.check_estimator(MyEstimator()) runs the library's conformance test suite, and parametrize_with_checks turns it into pytest tests. Declare special behaviour such as accepting NaN via __sklearn_tags__.

For a stateless function, a FunctionTransformer is simpler than a class.

Example

import numpy as np
from sklearn.base import BaseEstimator, OneToOneFeatureMixin, TransformerMixin
from sklearn.utils.validation import check_is_fitted, validate_data

class Winsorizer(OneToOneFeatureMixin, TransformerMixin, BaseEstimator):
    """Clip each feature to quantiles learned on the training data."""
    def __init__(self, lower=0.05, upper=0.95):
        self.lower = lower            # store params unchanged, nothing else
        self.upper = upper

    def fit(self, X, y=None):
        X = validate_data(self, X, reset=True)
        self.low_, self.high_ = np.quantile(X, [self.lower, self.upper], axis=0)
        return self

    def transform(self, X):
        check_is_fitted(self)
        X = validate_data(self, X, reset=False)
        return np.clip(X, self.low_, self.high_)

X = np.arange(100.0).reshape(-1, 1)
w = Winsorizer(upper=0.9).fit(X)
print(w.transform([[-50.0], [500.0]]).ravel())  # [ 4.95 89.1 ]
print(w.get_params())                           # {'lower': 0.05, 'upper': 0.9}

Remember

__init__ stores parameters unchanged; all learning happens in fit
Learned attributes end in _; fit returns self
Run check_estimator to verify the contract

Watch out for

Validating or transforming parameters in __init__, which breaks clone and set_params.

Check yourself

Why must __init__ store each argument unchanged under the same attribute name?

Official docs: scikit-learn.org/stable/developers/develop.html

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks