Tensorflow
Semantic Segmentation
Machine Learning
Deep Learning
Image Processing

Tensorflow How to ignore specific labels during semantic segmentation?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

In semantic segmentation, some pixels may be unlabeled, ambiguous, or intentionally excluded from training. The standard solution in TensorFlow is to mask those pixels out of the loss so they do not contribute to the gradient, instead of pretending they belong to a real class.

Ignore Label Is Not the Same as Background

This distinction is important:

  • background is a real class the model should learn
  • ignore label means "do not learn from this pixel"

If you map ignored pixels to background, the model learns the wrong thing. A masked loss is the correct approach.

Mask the Loss Explicitly

Assume your segmentation masks use 255 as the ignore label and valid class IDs range from 0 to num_classes - 1. You can compute sparse categorical cross-entropy only on valid pixels.

python
1import tensorflow as tf
2
3IGNORE_LABEL = 255
4
5def masked_sparse_ce(y_true, y_pred):
6    y_true = tf.cast(y_true, tf.int32)
7    valid_mask = tf.not_equal(y_true, IGNORE_LABEL)
8
9    safe_labels = tf.where(valid_mask, y_true, 0)
10
11    loss = tf.keras.losses.sparse_categorical_crossentropy(
12        safe_labels,
13        y_pred,
14        from_logits=True,
15    )
16
17    valid_mask = tf.cast(valid_mask, loss.dtype)
18    loss = loss * valid_mask
19
20    return tf.reduce_sum(loss) / tf.reduce_sum(valid_mask)

Then compile the model with that custom loss:

python
1model.compile(
2    optimizer="adam",
3    loss=masked_sparse_ce,
4    metrics=["accuracy"]
5)

The trick is replacing ignored labels temporarily with a safe dummy value before calling the loss, then zeroing those positions out with the mask.

Using sample_weight Instead

Keras also supports sample_weight, and for segmentation that weight can be a pixel-level mask. This is often a clean fit if your input pipeline already returns image, label, and weight tensors.

python
1def add_sample_weights(image, mask):
2    weights = tf.cast(tf.not_equal(mask, 255), tf.float32)
3    safe_mask = tf.where(mask == 255, 0, mask)
4    return image, safe_mask, weights

Now Model.fit can propagate the weights to the loss, so ignored pixels effectively count as zero.

This approach is especially nice when you also want class weighting. You can combine the ignore mask with per-class weights into one final pixel weight tensor.

Shape and Type Rules

For sparse segmentation losses, masks usually have shape (batch, height, width) or (batch, height, width, 1) and contain integer class IDs. Predictions usually have shape (batch, height, width, num_classes).

Make sure:

  • labels are integer typed
  • ignored pixels use a reserved label value
  • ignored label values are never passed directly into the sparse loss without masking

If you call sparse cross-entropy with out-of-range labels and no masking, TensorFlow will raise an error.

Metrics Need Attention Too

Masking only the loss is not always enough. If you report IoU or accuracy, ignored pixels should usually be excluded from those metrics as well. Otherwise the training objective and the reported evaluation numbers will disagree.

That means your metric function should apply the same valid_mask logic used by the loss.

Common Pitfalls

The biggest mistake is treating ignored pixels as background. That corrupts the learning target and often makes the model overpredict background.

Another mistake is masking the loss after reduction. If the loss has already been averaged across all pixels, it is too late to remove ignored positions correctly.

A third issue is forgetting the denominator. When all valid pixels are masked out in a batch item, you should guard against dividing by zero or make sure batching avoids fully empty supervision.

Summary

  • Ignore labels should be masked out of the loss, not remapped to a real class.
  • Background and ignore label are different concepts.
  • Use either a custom masked loss or pixel-level sample_weight.
  • Apply the same masking logic to metrics, not just training loss.
  • Validate label shapes and ranges early so ignored pixels do not trigger sparse-loss errors.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.