Join Our 5-Week ML/AI Engineer Interview Bootcamp 🚀 led by ML Tech Leads at FAANGs

Back to Questions

265. Image classification inference

easy
GeneralGeneral
senior

Implement a simple image classification inference step that turns a model’s raw scores (logits) into a predicted class and confidence for a single image. You’ll compute probabilities with softmax and return the top prediction in a beginner-friendly way.

The softmax formula is:

pi=ezi∑j=1Cezjp_i = \frac{e^{z_i}}{\sum_{j=1}^{C} e^{z_j}}

Requirements

Implement the function

python

Rules:

  • Compute a numerically stable softmax (subtract max(logits) before exponentiating).
  • Return the class with the highest probability and its probability as confidence.
  • Return probabilities as a NumPy array.
  • Do not use any prebuilt softmax helpers (e.g., scipy.special.softmax).
  • Use only NumPy.

Example

python

Output:

python
Input Signature
ArgumentType
logitsnp.ndarray
class_nameslist
Output Signature
Return NameType
valuetuple

Constraints

  • Use NumPy only.

  • Probabilities must be a NumPy array.

  • Use stable softmax.

Hint 1

Use np.exp(shifted_logits) where shifted_logits = logits - np.max(logits) for stability.

Hint 2

Normalize by dividing by the sum: probs = exp_vals / np.sum(exp_vals).

Hint 3

Use np.argmax(probs) to find the index of the highest probability.

Hint 4

Use logits = np.array(logits) if input is not already an array.

Roles
ML Engineer
AI Engineer
Companies
GeneralGeneral
Levels
senior
entry
Tags
softmax
numerical-stability
numpy
argmax
30 people are solving this problem
Python LogoPython Editor
Ln 1, Col 1

Input Arguments

Edit values below to test with custom inputs

You need tolog in/sign upto run or submit