Building A Custom Model in Scikit-Learn | by Tim Book | Towards Data Science

Scikit-Learn is incredible. It allows its users to fit almost any machine learning model you can think of, plus many you may never have even heard of! All in just two lines of code!
However, it doesn’t have everything. For example, ordinal regression is nowhere to be found. And its deep learning capabilities are... lacking. But who cares? You can find that stuff elsewhere, right?
True. But! Scikit-Learn isn’t just modeling. There are also some really great tools to help you streamline your modeling process, such as GridSearchCV and Pipeline. These tools are invaluable, but they only work with Scikit-Learn models. Turns out, if Scikit-Learn and our Google overlords don’t give it to us directly, we can make our own custom Scikit-Learn-compliant models! And it’s easier than you think!
In this post, I’m going to build something that is conspicuously missing from Scikit-Learn: the ability to use k-means clustering to do transfer learning in a Pipeline. That is, feeding the results of clustering into a supervised learning model in order to find the optimal value k.
This post is going to be a little technical. Specifically, I’m going to assume you have a working knowledge of object-oriented programming (OOP). That is, you know how and why to use the class Python keyword.
One of the best things about Scikit-Learn is its incredible consistency. Fitting one type of model is nominally the same as fitting any other type of model. That is, modeling in Scikit-Learn is as easy as:
model = MyModel(parameters)model.fit(X, y)
And that’s it! You can now analyze your model, probably with the help of the model’s .predict() and .score() methods. In fact, there are 5 methods every Scikit-Learn estimator is guaranteed to have:
.fit()
.predict()
.score()
.set_params()
.get_params()
To build our own model, we need only construct a class that has the above 5 methods and implements them in “the usual way.” That sounds like a lot of work. Luckily, Scikit-Learn does the hard work for us. In order to build our own model, we can inherit from a base class built into Scikit-Learn.
In OOP, if we specify that one class inherits from another, the “subclass” then gains all the methods of the “superclass.”
The syntax of inheritance looks like this:
class Car(Vehicle):    # stuff...
A car is a type of vehicle. The Car class will include every method of the Vehicle class, plus more that we can define in the Car class definition.
In order to comply with Scikit-Learn, our model needs to inherit from some mixin. A mixin is just a class that was never intended to work on its own, instead it just contains a lot of methods that you can add to your current class via inheritance. Scikit-Learn gives us one for each general type of model: RegressorMixin, ClassifierMixin, ClusterMixin, TransformerMixin, and several others we don’t need to worry about.
Everything we want to build ourselves, we make it by simply overriding what we inherit!
This has been a mouthful with no examples to show for it. First, a simple example. After, the motivation example using clustering.
The null model, sometimes called the “baseline” model, is the model where you have no information besides random guessing. For example, the null model for a regression problem would be just taking the mean y of your training data and using that as every prediction. For classification, it’s just taking the majority class for every prediction. For example, if you had to predict whether or not someone was going to win the lottery, the null model would dictate that you always lose, since that is the most likely outcome.
The null model is useful for telling how well your current model is doing. After all, if your model is good, it should beat the baseline. The null model is not built into Scikit-Learn, but it’s easy for us to implement!
import numpy as npfrom sklearn.base import RegressorMixinclass NullRegressor(RegressorMixin):    def fit(self, X=None, y=None):        # The prediction will always just be the mean of y        self.y_bar_ = np.mean(y)    def predict(self, X=None):        # Give back the mean of y, in the same        # length as the number of X observations        return np.ones(X.shape[0]) * self.y_bar_
Easy enough! We’re now free to do the usual...
model = NullRegressor()model.fit(X, y)model.predict(X)
The important part is, our new NullRegressor is now compatible with all of Scikit-Learn’s built-in tools such as cross_val_score and GridSearchCV.
This example was borne out of curiosity, when a coworker asked me if I could “tune” a k-means model using GridSearchCV and Pipeline. I originally said no, since you would need to use the clusterer as a transformer to pass into your supervised model, which Scikit-Learn doesn’t allow for. But why let that stop us? We just learned how to hack Scikit-Learn to do whatever we want! To spell it out, essentially what I want is to create the following pipeline:
Pipeline([    ("sc", StandardScaler()),    ("km", KMeansSomehow()),    ("lr", LogisticRegression()])
Where KMeansSomehow() is a clusterer used as a Scikit-Learn transformer. That is, it appends the onehot-encoded cluster labels to the data matrix X to then pass into our model. In order for this to work, we’ll start by defining a class that inherits from TransformerMixin. We’ll then give it the appropriate .fit(), .transform(), and .fit_transform() methods.
But to start, the initialization:
from sklearn.base import TransformerMixinfrom sklearn.cluster import KMeansclass KMeansTransformer(TransformerMixin):    def __init__(self, *args, **args):        self.model = KMeans(*args, **args)
The purpose of self.model is to contain the underlying cluster model. But what are *args and **kwargs, you ask? They’re a shortcut for lazy programmers. They essentially capture all other arguments you pass into __init__() and pass them right along to KMeans(). This is essentially me saying “Whatever I pass into KMeansTransformer will also pass into KMeans, and I’m too lazy to figure out what those arguments might be in the future.”
Next, we need to give it the appropriate fitting methods:
from self.preprocessing import OneHotEncoderclass KMeansTransformer(TransformerMixin):    def __init__(self, *args, **args):        self.model = KMeans(*args, **args)    def fit(self, X):        self.X = X        self.model.fit(X)    def transform(self, X):        # Need to reshape into a column vector in order to use        # the onehot encoder.        cl = self.model.predict(X).reshape(-1, 1)                self.oh = OneHotEncoder(            categories="auto",             sparse=False,            drop="first"        )        cl_matrix = self.oh.fit_transform(cl)               return np.hstack([self.X, cl_matrix])    def fit_transform(self, X, y=None):        self.fit(X)        return self.transform(X)
And that should be it! We can now use this KmeansTransformer just like a built-in Scikit-Learn transformer. To top it off, try this example out:
from sklearn.preprocessing import StandardScalerfrom sklearn.linear_model import LogisticRegressionfrom sklearn.datasets import make_blobsX, y = make_blobs(    n_samples=100,    n_features=2,    centers=3)pipe = Pipeline([    ("sc", StandardScaler()),    ("km", KMeansTransformer()),    ("lr", LogisticRegression(penalty="none", solver="lbfgs"))])pipe.fit(X, y)pipe.score(X, y)# ==> 1.0
In less artificial examples, you might also want to use GridSearchCV to find the optimal number of clusters to pass into your logistic regression (or whatever model you have).
Now you should understand how to build your own custom machine learning models within the framework of Scikit-Learn, which is currently the most popular and (in many cases) powerful ML library out there.
Is this blog post pedantic? Sure. Will you ever need this? No, you probably won’t ever need this, but that’s not the point. This technique is used to build something. It takes a builder to recognize the need to build a more creative solution. Plus, as a wise person once put it:
You are never less for knowing more.
I hope you appreciate adding this tool to your tool belt. Master it, and you may just achieve true Machine Learning Ascension.