The `Classifier` class

The Base Classification Model

The shared classifier base

A small Classifier base class that every classification model in the book inherits from. Same role as d2l.Module for regression — but with classification-specific defaults:

A validation step that reports loss and accuracy.
An accuracy helper that compares the argmax of the predicted scores to the true labels.

Subclasses just supply forward (and a custom loss if not plain cross-entropy).

Scores, probabilities, decisions

Classifiers usually produce a vector of scores \mathbf{o}\in\mathbb{R}^q. The training loss may turn scores into probabilities, but the deployed decision is often just

\hat{y}=\arg\max_j o_j.

Keep the roles separate:

scores/logits: differentiable quantities the model outputs;
loss: smooth training signal, e.g. cross-entropy;
accuracy: discrete evaluation metric after taking argmax.

Accuracy is what many benchmarks report, but it is not a useful gradient: one tiny score change usually leaves argmax unchanged.

Base classifier imports

from d2l import jax as d2l
from functools import partial
from jax import numpy as jnp
import jax
import optax

class Classifier(d2l.Module):
    """The base class of classification models."""
    def training_step(self, params, batch, state):
        # Here value is a tuple since models with BatchNorm layers require
        # the loss to return auxiliary data
        value, grads = jax.value_and_grad(
            self.loss, has_aux=True)(params, batch[:-1], batch[-1], state)
        l, _ = value
        self.plot("loss", l, train=True)
        return value, grads

    def validation_step(self, params, batch, state):
        # Discard the second returned value. It is used for training models
        # with BatchNorm layers since loss also returns auxiliary data
        l, _ = self.loss(params, batch[:-1], batch[-1], state)
        self.plot('loss', l, train=False)
        self.plot('acc', self.accuracy(params, batch[:-1], batch[-1], state),
                  train=False)

Default optimizer hook

A default configure_optimizers on Module so subclasses don’t have to write it:

@d2l.add_to_class(d2l.Module)
def configure_optimizers(self):
    return optax.sgd(self.lr)

Accuracy

Take the argmax along the class axis, compare with the true label element-wise, and average. The result is the fraction of correctly-classified examples in the batch:

@d2l.add_to_class(Classifier)
@partial(jax.jit, static_argnums=(0, 5))
def accuracy(self, params, X, Y, state, averaged=True):
    """Compute the number of correct predictions."""
    Y_hat = state.apply_fn({'params': params,
                            'batch_stats': state.batch_stats},  # BatchNorm Only
                           *X)
    Y_hat = d2l.reshape(Y_hat, (-1, Y_hat.shape[-1]))
    preds = d2l.astype(d2l.argmax(Y_hat, axis=1), Y.dtype)
    compare = d2l.astype(preds == d2l.reshape(Y, (-1,)), d2l.float32)
    return d2l.reduce_mean(compare) if averaged else compare

The validation step then reports both the loss (lower is better) and accuracy (higher is better) every epoch.

Why report both loss and accuracy?

Two models can have the same accuracy but different confidence. Cross-entropy still notices whether the correct class received probability 0.51 or 0.99.

Use both during training:

loss detects calibration and optimization progress;
accuracy tracks the hard decision quality students and benchmarks usually care about;
disagreement between them is diagnostic, not a bug.

Recap

Classifier(d2l.Module) adds accuracy reporting to the base scaffold from the regression chapter.
One line for accuracy: argmax → ==y → mean.
The same training loop now drives every classification model we’ll build through the rest of the book.

The Classifier class