Regression vs classification output

Background

Vision models fall into two output families:

• Regression — a continuous value (age, depth, a motion vector). The model output is the answer.
• Classification — a discrete label out of $K$ classes. The answer is the index of the highest-scoring class.

Problem statement

Implement vision_task_output(z, task) that returns the right prediction for the task:

• task == "regression" → return the continuous value z as a float.
• task == "classification" → return the predicted class index (argmax of the score vector z) as an int.

Input

• z — for regression, a scalar; for classification, an array-like of K scores/logits.
• task — "regression" or "classification".

Output

A float (regression) or an int class index (classification).

Examples

Example 1 — regression

Input:  z = 3.7, task = "regression"
Output: 3.7

Example 2 — classification

Input:  z = [1, 3, 2], task = "classification"
Output: 1

Explanation: index 1 has the largest score.

Constraints

• Regression returns float(z).
• Classification returns int(argmax(z)).

Notes

• The choice of head (and its loss — squared error for regression, cross-entropy for classification) is what specialises the same backbone to different vision tasks.