PyTorch: logits, class labels and cross-entropy

pythontorchindependent-practice

Saved on this device

Full lesson · 4 learning steps Open any exercise for focused practice

Step 1 of 4 · intro card

PyTorch: logits, class labels and cross-entropy

Concept and contract

Match output and target contracts before interpreting a loss value.

A classification loss connects model outputs with target semantics. In the common class-index form, CrossEntropyLoss accepts unnormalized logits shaped batch by classes and integer class indices shaped batch.

It applies the required log-softmax internally. Feeding already normalized probabilities changes the intended calculation, and floating class indices do not match this particular target contract.

Mechanism and worked example

The helper creates float32 score tensors and long integer labels, checks the batch axes and computes cross-entropy directly from logits.

The helper creates float32 score tensors and long integer labels, checks the batch axes and computes cross-entropy directly from logits. Argmax over dimension one yields one predicted class per example.

Equal two-class logits produce a loss near log two, which gives a simple numerical baseline. Prediction confidence and correctness remain separate from that scalar loss.

Complete worked example

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def classification_loss(logits, labels):
    scores = torch.tensor(logits, dtype=torch.float32)
    target = torch.tensor(labels, dtype=torch.long)
    if scores.ndim != 2 or target.ndim != 1 or scores.shape[0] != target.shape[0]:
        raise ValueError('batch and class axes required')
    loss = nn.CrossEntropyLoss()(scores, target)
    return {'loss': loss.item(), 'predictions': scores.argmax(dim=1).tolist()}

Reasoning: The helper creates float32 score tensors and long integer labels, checks the batch axes and computes cross-entropy directly from logits. Argmax over dimension one yields one predicted class per example.

Counterexamples and common mistakes

The counterexample sends softmax probabilities into CrossEntropyLoss, effectively normalizing twice.

The counterexample sends softmax probabilities into CrossEntropyLoss, effectively normalizing twice. Argmax predictions may remain unchanged while the loss becomes wrong.

A prediction-only test would miss the defect. Also avoid reducing over the batch axis when selecting classes, and validate label ranges when accepting external data rather than relying solely on a low-level indexing exception.

Counterexample: a runnable framework mistake

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def classification_loss(logits, labels):
    scores = torch.tensor(logits, dtype=torch.float32)
    target = torch.tensor(labels, dtype=torch.long)
    if scores.ndim != 2 or target.ndim != 1 or scores.shape[0] != target.shape[0]:
        raise ValueError('batch and class axes required')
    loss = nn.CrossEntropyLoss()(scores.softmax(dim=1), target)
    return {'loss': loss.item(), 'predictions': scores.argmax(dim=1).tolist()}

Reasoning: The counterexample sends softmax probabilities into CrossEntropyLoss, effectively normalizing twice. Argmax predictions may remain unchanged while the loss becomes wrong.

Independent practice: one bounded framework change

Add an accuracy field by comparing the predicted class indices with target and taking their floating mean.

Add an accuracy field by comparing the predicted class indices with target and taking their floating mean. Preserve the raw-logit loss calculation and prediction list.

This small metric addition should not modify the tensors or insert a softmax. Keep the dimensionality contract visible and use the same batch population for both loss and accuracy.

Step 1: Add accuracy over the same example batch without changing the loss input. Step 2: Compare argmax class indices with target.

Practice starter: solve before reading the solution

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def classification_loss(logits, labels):
    # Remove this placeholder, then adapt the retained baseline:
    # Add accuracy over the same example batch without changing the loss input.
    raise NotImplementedError
    scores = torch.tensor(logits, dtype=torch.float32)
    target = torch.tensor(labels, dtype=torch.long)
    if scores.ndim != 2 or target.ndim != 1 or scores.shape[0] != target.shape[0]:
        raise ValueError('batch and class axes required')
    loss = nn.CrossEntropyLoss()(scores, target)
    return {'loss': loss.item(), 'predictions': scores.argmax(dim=1).tolist()}

Reasoning: Add accuracy over the same example batch without changing the loss input. Compare argmax class indices with target.

Hint 1: Convert the boolean comparison to float before mean.
Hint 2: Keep raw logits for CrossEntropyLoss.

Solution, tests and trade-offs

Fixtures include equal logits, confident correct examples and a mismatched batch.

Fixtures include equal logits, confident correct examples and a mismatched batch. Practice adds mixed correctness and compares finite loss and accuracy values with tolerances.

A lower loss on these tiny fixed inputs does not demonstrate model generalization. The cases specifically separate numerical loss semantics, class-axis selection and agreement with the supplied target indices.

Practice solution and acceptance checks

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def classification_loss(logits, labels):
    scores = torch.tensor(logits, dtype=torch.float32)
    target = torch.tensor(labels, dtype=torch.long)
    if scores.ndim != 2 or target.ndim != 1 or scores.shape[0] != target.shape[0]:
        raise ValueError('batch and class axes required')
    loss = nn.CrossEntropyLoss()(scores, target)
    accuracy = (scores.argmax(dim=1) == target).float().mean().item()
    return {'accuracy': accuracy, 'loss': loss.item(), 'predictions': scores.argmax(dim=1).tolist()}

Reasoning: Fixtures include equal logits, confident correct examples and a mismatched batch. Practice adds mixed correctness and compares finite loss and accuracy values with tolerances.

Interview reasoning and next step

Other tasks need different losses: independent binary labels often use BCEWithLogitsLoss, while regression may use MSELoss or another domain-appropriate objective..

Other tasks need different losses: independent binary labels often use BCEWithLogitsLoss, while regression may use MSELoss or another domain-appropriate objective.

Choose the target representation before writing the output layer. Document whether metrics average over examples, classes or tokens, and avoid applying an activation merely because it makes displayed outputs look like probabilities.

Optional stretch — after the core task works: Compare macro recall with accuracy on an imbalanced held-out set.

Computer practice: https://applydjinn.com/typorithm/python-frameworks-lab-20261105.zip. Install the pinned dependencies in an isolated environment, run verify.py --track torch, then edit the retained exercise baseline and check it with --practice pt-losses. The lab uses owned in-memory databases, local framework clients and tiny CPU tensors. It never reads the working ApplyDjinn database, contacts production feeds or downloads models.

Match output and target contracts before interpreting a loss value.