Introduction

In Keras, class_weight works for standard classification targets, but it does not handle dense pixel-wise segmentation targets such as those used in U-Net. For semantic segmentation, the usual solution is to build the class weights into the loss function itself so each pixel contributes according to its class. That gives you the same effect as class weighting, but at the per-pixel level.

Why class_weight Fails For U-Net Targets

A U-Net model typically predicts a class for every pixel, so the target tensor has height and width dimensions in addition to the batch dimension. Keras's class_weight argument is not designed for that shape.

Instead of trying to push class weighting through model.fit(..., class_weight=...), define a weighted segmentation loss.

Weighted Sparse Categorical Crossentropy

If your masks store one integer class ID per pixel, weighted sparse categorical crossentropy is a good choice.

1import tensorflow as tf
2
3class_weights = tf.constant([0.2, 2.0, 4.0], dtype=tf.float32)
4
5
6def weighted_sparse_cce(y_true, y_pred):
7    y_true = tf.cast(y_true, tf.int32)
8    pixel_weights = tf.gather(class_weights, y_true)
9
10    loss = tf.keras.losses.sparse_categorical_crossentropy(
11        y_true, y_pred
12    )
13    return tf.reduce_mean(loss * pixel_weights)

Here:

• 'y_true contains class indices such as 0, 1, or 2'
• 'y_pred contains softmax probabilities for each class'
• 'tf.gather maps each pixel label to its class weight'

That is the core pattern most segmentation weighting strategies build on.

Minimal U-Net-Like Example

The following example compiles a small segmentation model using the custom loss.

1import tensorflow as tf
2
3inputs = tf.keras.Input(shape=(64, 64, 1))
4x = tf.keras.layers.Conv2D(16, 3, padding="same", activation="relu")(inputs)
5x = tf.keras.layers.Conv2D(16, 3, padding="same", activation="relu")(x)
6outputs = tf.keras.layers.Conv2D(3, 1, activation="softmax")(x)
7
8model = tf.keras.Model(inputs, outputs)
9model.compile(optimizer="adam", loss=weighted_sparse_cce)
10
11x_batch = tf.random.normal((2, 64, 64, 1))
12y_batch = tf.random.uniform((2, 64, 64), maxval=3, dtype=tf.int32)
13
14model.train_on_batch(x_batch, y_batch)

This is not a full U-Net architecture, but it demonstrates the loss integration clearly.

One-Hot Masks Need A Slightly Different Loss

If your masks are one-hot encoded instead of integer encoded, use categorical crossentropy and derive the class ID with argmax for weighting.

1import tensorflow as tf
2
3class_weights = tf.constant([0.2, 2.0, 4.0], dtype=tf.float32)
4
5
6def weighted_cce(y_true, y_pred):
7    class_ids = tf.argmax(y_true, axis=-1, output_type=tf.int32)
8    pixel_weights = tf.gather(class_weights, class_ids)
9
10    loss = tf.keras.losses.categorical_crossentropy(y_true, y_pred)
11    return tf.reduce_mean(loss * pixel_weights)

Choose the version that matches how your masks are stored.

Choosing The Weights

Class weights often come from inverse frequency or median-frequency balancing. For example, if background pixels dominate the dataset, background usually gets a smaller weight and rare foreground classes get larger weights.

The exact values are empirical. Extremely large weights can destabilize training, so start conservatively and monitor both loss and segmentation quality metrics such as Dice score or IoU.

Common Pitfalls

A common mistake is trying to force class_weight into a segmentation problem with dense target tensors. Keras does not apply it the way people expect for pixel-wise masks.

Another mistake is mixing sparse masks with categorical loss, or one-hot masks with sparse loss. The loss function must match the target representation.

Developers also often forget that the final model layer and loss must agree. softmax plus categorical-style loss is the usual multi-class setup.

Finally, weighting alone does not solve every imbalance problem. Patch sampling, focal loss, Dice-based losses, or hybrid losses may still be necessary when small structures are hard to learn.

Summary

• 'class_weight is not the right tool for dense segmentation masks in U-Net training.'
• Use a custom per-pixel weighted loss instead.
• For integer masks, weighted sparse categorical crossentropy is a good default.
• For one-hot masks, use a weighted categorical crossentropy variant.
• Keep the target encoding, output activation, and loss function aligned.
• Tune class weights carefully and evaluate with segmentation-specific metrics, not loss alone.