Implementing Custom Estimators

Implementing your own estimator requires very little code, which is why most scikit-learn users build up a collection of custom models over time.

Programming

From Custom Transformers to Custom Models

Once you have seen how scikit-learn lets you write your own transformer, a natural next question is whether the same approach works for models that make predictions. It does. Creating a custom classifier or regressor follows almost the same path as creating a custom transformer. Only two details change, and once you know them you can read the pattern in either direction.

As you work through this article, keep one question in mind: if you already know how a custom transformer is built, what exactly would you swap out to build a custom classifier or regressor, and what stays the same?

Which Base Class Changes

A custom transformer inherits from TransformerMixin. In the classifier and regressor pattern, that one parent is swapped out. A custom classifier inherits from ClassifierMixin instead, and a custom regressor inherits from RegressorMixin instead. The mixin is the base class you inherit from, and it supplies the behavior that belongs to that kind of estimator.

Which parent classes does a custom classifier inherit from, and which parent supplies which behavior? Both custom models share BaseEstimator, but each one picks a different mixin.

Read the diagram from the bottom up to see the design. Your class is the terminal node at the bottom. Each arrow shows an inheritance link, meaning your class receives the behavior of the parent it points to. The only difference between a classifier and a regressor in this structure is which mixin sits on the second level.

Predict Replaces Transform

The second change is the method you write. A transformer implements transform. A classifier or regressor implements predict instead. Everything else about the structure follows the transformer pattern: you choose the right base class and then write the one method that does the work of producing outputs for new data.

How does the method a custom classifier implements differ from the transform method a transformer implements? Read each column from top to bottom.

The two columns differ only in the names of their boxes. The structure is the same: a base class first, then one method you implement. The changed entries are the base class and the method, which is exactly what the source identifies as the two details that change.

Tracing a Minimal Custom Classifier

The following classifier is a small illustrative example written for this article. It predicts the most common label seen during training for every input. It is not from the source and is included only to show how the two changes look in actual code. Read it as a trace: the base classes are on the class line, the training step stores state, and predict uses that stored state.

python

from sklearn.base import BaseEstimator, ClassifierMixin
import numpy as np

class MajorityClassifier(BaseEstimator, ClassifierMixin):
    def fit(self, X, y):
        values, counts = np.unique(y, return_counts=True)
        self.classes_ = values
        self.majority_ = values[np.argmax(counts)]
        return self

    def predict(self, X):
        return np.full(len(X), self.majority_)

from sklearn.base import BaseEstimator, ClassifierMixin
import numpy as np

class MajorityClassifier(BaseEstimator, ClassifierMixin):
    def fit(self, X, y):
        values, counts = np.unique(y, return_counts=True)
        self.classes_ = values
        self.majority_ = values[np.argmax(counts)]
        return self

    def predict(self, X):
        return np.full(len(X), self.majority_)

Walk through it

1
Line 1 imports BaseEstimator and ClassifierMixin from sklearn.base
2
Line 4 declares the class and lists ClassifierMixin as a parent
3
fit counts how often each label appears in y and stores the most frequent label as majority_
4
fit returns self so the fitted object can be used directly
5
predict returns one copy of the stored majority label for each row in X

Notes on the code

The class inherits from ClassifierMixin, the classifier mixin named in the source
predict is the method implemented in place of transform
fit stores the learned state on the object during training

Worked example

Tracing the state of MajorityClassifier

Train MajorityClassifier on the labels y = [0, 1, 1, 1, 0], then predict for three new rows.

1
Count the labels
The label 0 appears 2 times and the label 1 appears 3 times.
2
Store the fitted state
fit sets classes_ to [0, 1] and majority_ to 1, because 1 has the highest count.
3
Predict on new data
predict ignores the values inside X and returns the stored label 1 once for each of the three rows.

Result

The classifier returns [1, 1, 1]. The fitted state majority_ = 1 is what flows from fit into predict.

Why Custom Estimators Stay Small

The source makes a point about the size of the job. Implementing your own estimator requires very little code. Because the inheritance and method names are the main changes from the transformer pattern, most of the work is choosing the right base class and writing the one method that does the prediction work. The rest of the structure is already supplied by the mixin you inherit from.

This small cost is the reason most scikit-learn users build up a collection of custom models over time. Once the pattern is familiar, adding another model to that collection takes only a short class with the right base class and a predict method.

Mistakes When Switching From Transformers

Common mistake
Inheriting from TransformerMixin in a classifier
What it looks like
```
class MyModel(BaseEstimator, TransformerMixin) followed by a predict method
```
Why it fails
TransformerMixin belongs to the transformer pattern. A classifier inherits from ClassifierMixin.
Do this instead
Replace TransformerMixin with ClassifierMixin, or RegressorMixin for a regressor.
Common mistake
Writing transform in a classifier or regressor
Why it fails
The source says a classifier or regressor implements predict instead of transform.
Do this instead
Rename the method to predict and make it return predictions for the input rows.
Common mistake
Keeping the transformer base class and the classifier method together, so the class mixes the two patterns

Checking Your Understanding of the Method Swap

What do you think happens?

You are converting a custom transformer into a custom regressor. Which method do you write in place of transform?

transform
predict
fit_transform
score

Reveal answer

Answer:

predict

The source states that a classifier or regressor implements predict instead of transform.

EASY

A colleague writes a custom regressor that inherits from TransformerMixin and implements transform. Name the two changes you would make to turn it into a correct custom regressor.

Hints

Look at the two details the source says change between the transformer pattern and the classifier or regressor pattern.
One change is in the base class and one change is in the method name.

Your attempt (not saved)

Key Takeaways

A custom classifier inherits from ClassifierMixin and a custom regressor inherits from RegressorMixin, in place of TransformerMixin.
A custom classifier or regressor implements predict in place of transform, and the rest of the pattern stays the same as the transformer pattern.
Because only the base class and the method name change, a custom estimator requires very little code.
That small cost is why most scikit-learn users build up a collection of custom models over time.

More concepts starting with I