Difference Between predict and predict_proba in scikit-learn | Towards Data Science

When training models (and more precisely supervised estimators) with sklearn, we sometimes need to predict the actual class while in some other occasions we may want to predict the class probabilities.
In today’s article we will discuss how to use predict and predict_proba methods over a dataset in order to perform predictions. Additionally, we’ll explore the differences between these methods and discuss when to use one over the other.
First, let’s create an example model that we’ll reference throughout this article in order to demonstrate a few concepts. In our examples, we will be using the Iris dataset which is also included in sklearn.datasets module of scikit-learn. This will be a classification task in which we need to identify and correctly predict three distinct types of irises, namely Setosa, Versicolour, and Virginica from the petal and sepal dimensions (length and width).
import numpy as npfrom sklearn import datasetsfrom sklearn.neighbors import KNeighborsClassifier# Load the Iris datasetiris_X, iris_y = datasets.load_iris(return_X_y=True)# Split Iris dataset into train/test sets randomlynp.random.seed(0)indices = np.random.permutation(len(iris_X))iris_X_train = iris_X[indices[:-10]]iris_y_train = iris_y[indices[:-10]]iris_X_test = iris_X[indices[-10:]]iris_y_test = iris_y[indices[-10:]]# Instantiate and fit a KNeighbors classifierknn = KNeighborsClassifier()knn.fit(iris_X_train, iris_y_train)
All supervised estimators in scikit-learn implement the predict() method that can be executed on a trained model in order to predict the actual label (or class) over a new set of data.
The method accepts a single argument that corresponds to the data over which the predictions will be made and it returns an array containing the predicted label for each data point.
predictions = knn.predict(iris_X_test)print(predictions)array([1, 2, 1, 0, 0, 0, 2, 1, 2, 0])
In the context of classification tasks, some sklearn estimators also implement the predict_proba method that returns the class probabilities for each data point.
The method accepts a single argument that corresponds to the data over which the probabilities will be computed and returns an array of lists containing the class probabilities for the input data points.
predictions = knn.predict_proba(iris_X_test)print(predictions)array([[0. , 1. , 0. ],       [0. , 0.4, 0.6],       [0. , 1. , 0. ],       [1. , 0. , 0. ],       [1. , 0. , 0. ],       [1. , 0. , 0. ],       [0. , 0. , 1. ],       [0. , 1. , 0. ],       [0. , 0. , 1. ],       [1. , 0. , 0. ]])
In today’s article we discussed how to perform predictions over data using a pre-trained scikit-learn model. Additionally, we explored the main differences between the methods predict and predict_proba which are implemented by estimators of scikit-learn.
The predict method is used to predict the actual class while predict_proba method can be used to infer the class probabilities (i.e. the probability that a particular data point falls into the underlying classes).
Become a member and read every story on Medium. Your membership fee directly supports me and other writers you read.
You may also like