Semantic Segmentation Fundamentals

Build the mental model

A classification model outputs a single class-score vector (shape: num_classes) for the entire image, answering one question: 'what is in this image?' Semantic segmentation answers a much broader question — 'which class does each individual pixel belong to?' Because of this, a segmentation model's output is not a single vector but a tensor of shape (batch, num_classes, H, W) that preserves the image's spatial dimensions — you can think of it as an independent set of class scores at every single pixel position.

To get the final predicted class at each pixel, you take argmax along the class dimension (dim=1) — collapsing the num_classes dimension down to just the index of the highest-scoring class at that position. The result is an integer tensor of shape (batch, H, W), where the value at position (i, j) is the predicted class index for that pixel. Visualized, this output looks like a mask that partitions the whole image into class-colored regions — a complete contrast to classification's single label per image.

Connect it to a real scenario

Suppose the Tutorial Platform is considering a feature that segments user-uploaded lesson diagrams into 'text region' versus 'illustration region' to generate better alt-text for screen-reader users. Only with a per-pixel segmentation output like this can you route the text region to an OCR pipeline and the illustration region to a separate visual-description pipeline.

Classification produces a single label per image, while segmentation produces a label per pixel.
A segmentation model's output shape is (batch, num_classes, H, W), preserving the spatial dimensions.
Applying argmax(dim=1) over the class dimension yields a per-pixel class map of shape (batch, H, W).

Try the working example

python

import torch

batch_size, num_classes, height, width = 2, 5, 8, 8

# Fake per-pixel class scores from a segmentation model
segmentation_output = torch.randn(batch_size, num_classes, height, width)

# Collapse the class dimension to get the predicted class per pixel
predicted_mask = segmentation_output.argmax(dim=1)

print(f"segmentation_output shape: {tuple(segmentation_output.shape)}")
print(f"predicted_mask shape: {tuple(predicted_mask.shape)}")
print(f"predicted_mask dtype: {predicted_mask.dtype}")

You should see

Prints exactly: 'segmentation_output shape: (2, 5, 8, 8)', 'predicted_mask shape: (2, 8, 8)', 'predicted_mask dtype: torch.int64' — argmax removes the class dimension (dim=1, size 5) and returns an integer index tensor.

5-minute try-it

Count how many pixels in predicted_mask were assigned class index 2 using (predicted_mask == 2).sum() and print the result, then change num_classes from 5 to 20 and observe how the shapes change.

One important caution

Using argmax(dim=0) instead of argmax(dim=1) collapses the batch dimension instead of the class dimension, producing a meaningless (num_classes, H, W) output rather than a per-pixel class map.

Feeding predicted_mask (dtype int64) directly while leaving a ground-truth mask in float dtype causes nn.CrossEntropyLoss to raise a dtype mismatch error, since it expects integer class-index targets.

Wikipedia — Image segmentation — Computer Vision