How can I visualize border/decision function of two classes using scikit-learn
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.
Introduction
Visualizing a binary classifier’s decision boundary is one of the quickest ways to understand what the model has learned. In scikit-learn, the standard pattern is to train on two features, build a mesh grid across that 2D space, score the grid, and then plot either predicted classes, probabilities, or the raw decision function.
Train a Model on Two Features
Decision-boundary plots are easiest when the training data has exactly two input features. A small synthetic dataset keeps the example focused on the plotting pattern.
Using a pipeline is important because the grid points you plot will go through the same preprocessing as the training data.
Build the Mesh Grid
The next step is to create a dense grid of coordinates covering the region where your data lives.
Each row in grid is a synthetic point in feature space. Once the classifier predicts on those points, you can reshape the result back to the mesh and plot it.
Plot Predicted Class Regions
The simplest picture is a region plot showing which class the model predicts at each location.
This is a good first view, but it hides how strongly the classifier prefers one class over the other near the boundary.
Plot the Decision Function
For models that expose decision_function, plotting the score surface is usually more informative. The contour where the score equals zero is the decision boundary.
This view tells you more than class labels alone. It shows how sharply the model separates the two classes and where the margin is narrow.
Plot Probabilities When Available
If the estimator supports predict_proba, you can visualize the probability surface instead of raw decision scores.
This is useful when your discussion is about confidence or thresholding rather than just classification regions.
What Changes with Different Models
The plotting workflow stays almost the same across classifiers, but the geometry changes:
- linear models produce straight boundaries in feature space,
- kernel methods can produce curved shapes,
- decision trees create axis-aligned regions,
- ensemble models can create more jagged surfaces.
That means the plot is not just decoration. It is a direct picture of the model family you chose.
What If the Original Data Has More Than Two Features
You cannot draw the true boundary of a high-dimensional classifier on a flat 2D chart. In practice, people use one of these strategies:
- train a model on only two selected features,
- project the data to two dimensions for visualization,
- hold the other features constant and plot a slice.
All of these are approximations. The plot can still be useful, but it is important to say that it is a projection or slice rather than the full decision geometry.
Common Pitfalls
- Forgetting to apply the same preprocessing to the mesh grid that the model used during training.
- Using a very coarse grid and getting a blocky, misleading boundary.
- Plotting hard class labels only and missing the confidence structure around the separating line.
- Treating a 2D projection as if it were the full boundary of a high-dimensional model.
- Assuming
predict_probaanddecision_functionmean the same thing for every estimator.
Summary
- Train on two features or create a deliberate 2D slice before plotting a decision boundary.
- Build a dense mesh grid, score it with the model, and reshape the output for plotting.
- Use
predictfor class regions,decision_functionfor margin-like scores, andpredict_probafor probability surfaces. - Keep preprocessing inside a pipeline so the plotted boundary matches the trained model.
- Be explicit when the visualization is only a slice or projection of a higher-dimensional problem.
Related reading
- How can I visualize the weightsvariables in cnn in Tensorflow?
- How can implement subsample like keras in tensorflow?
- How can neural networks learn functions with a variable number of inputs?
- How can Tensorboard files be merged/combined or appended?
- How can I write a table function to coerce the order of column names in the output of table?
- How can the Euclidean distance be calculated with NumPy?
- How can use regularization in TensorFlow-Slim?
- how can you train a model twice multiple times in sklearn using fit.?
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack what you have practised
A free account saves your progress, solutions and study plan across every problem on Codemia.
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.