argmax function used in Machine Learning in Python

In this tutorial, we will be learning about the Argmax function used in machine learning.

The argmax (arguments of the maxima) in mathematics, are points of a function at which it is maximized. Arg max to the arguments at which function archives its max value. Whereas global maxima generally refers to the largest possible outputs of a function,

Consider an example where f(x) is calculated as the cube of the x, input values (x) is limited to integers from 1 to 5:

  • f(1) = 1^3 = 1
  • f(2) = 2^3 = 8
  • f(3) = 3^3 = 27
  • f(4) = 4^3 = 64
  • f(5) = 5^3 = 125

We can intuitively see that the argmax for the function f(x) is 5.
Note that this function neither returns the maximum value of the function nor it returns the argument having maximum value, rather it returns the arguments at which the function archives its maxima.

Argmax function in Machine Learning

The function is used in various ways throughout the field of machine learning calculations.
Let me take an example of a multi-class classification: 1st”, “2nd”, “3rd”, and “4th” are the four classes. The mapping is as follows:

  1. First-class = 0
  2. Second-class = 1
  3. Third-class = 2
  4. Fourth-class = 3

A model has made one prediction for an input sample and predicted the following vector of probabilities:

  • vect = [0.2, 0.3, 0.4, 0.1]

We can see that the example has a 20 percent likelihood of belonging to 1st class, 30 percent to 2nd, and 40 percent to 3rd, 10 percent to 4th.

  • argmax (vect) = 2, as at index the maximum value is achieved.

Implementation of NumPy argmax function using Python

Firstly, using the inbuilt argmax() function present in the NumPy library.

from numpy import argmax

vect = [0.2, 0.3, 0.4, 0.1]

ans = argmax(vect)
print('arg max of %s: %d' % (vect, ans))


arg max of [0.2, 0.3, 0.4, 0.1]: 2

Let’s now perform the same task as above, but this time we will implement it using a user-defined function.
Define a function that would take a vector as an argument and return the index corresponding to the maximum value.

def argmax(vect):
  index, val = 0, vect[0]
  for i,v in enumerate(vect):
    if v > val:
      index, val = i,v
  return index

vect = [0.2, 0.3, 0.4, 0.1]

ans = argmax(vect)
print('arg max of %s: %d' % (vect, ans))


arg max of [0.2, 0.3, 0.4, 0.1]: 2

In a realistic scenario, one shall deal with a collection of multiple samples. So let’s now store the values in a matrix form. Taking a matrix with each column representing a class label, rows of predicted probabilities.

from numpy import argmax
from numpy import asarray

prob = asarray([[0.4, 0.2, 0.4], [0.0, 1.0, 0.0], [0.7, 0.2, 0.1], [0.3, 0.3, 0.4]])

ans = argmax(prob, axis=1)


(4, 3)
[0 1 0 2]

Hope you enjoyed learning with me on this topic. I would request you to practice the application of this function on various other samples. Thank you and happy learning.

Also read: Predict Weather Report Using Machine Learning in Python

Leave a Reply

Your email address will not be published. Required fields are marked *