Tensorflow: How to ignore specific labels during semantic segmentation?

Question

I'm using tensorflow for semantic segmentation. How can I tell tensorflow to ignore a specific label when computing the pixelwise loss?

I've read in this post that for image classification one can set the label to -1 and it will be ignored. If that is true, given the label-tensor, how can I modify my labels such that certain values are changed to -1?

In Matlab it would be something like:

ignore_label = 255
myLabelTensor(myLabelTensor == ignore_label) = -1

But I don't know how to do this in TF?

Some background info:
This is how the labels are loaded:

label_contents = tf.read_file(input_queue[1])
label = tf.image.decode_png(label_contents, channels=1)

This is how the loss is currently computed:

raw_output = net.layers['fc1_voc12']
prediction = tf.reshape(raw_output, [-1, n_classes])
label_proc = prepare_label(label_batch, tf.pack(raw_output.get_shape()[1:3]),n_classes)
gt = tf.reshape(label_proc, [-1, n_classes])

# Pixel-wise softmax loss.
loss = tf.nn.softmax_cross_entropy_with_logits(prediction, gt)
reduced_loss = tf.reduce_mean(loss)

with

def prepare_label(input_batch, new_size, n_classes):
    """Resize masks and perform one-hot encoding.

    Args:
      input_batch: input tensor of shape [batch_size H W 1].
      new_size: a tensor with new height and width.

    Returns:
      Outputs a tensor of shape [batch_size h w 21]
      with last dimension comprised of 0's and 1's only.
    """
    with tf.name_scope('label_encode'):
        input_batch = tf.image.resize_nearest_neighbor(input_batch, new_size) # as labels are integer numbers, need to use NN interp.
        input_batch = tf.squeeze(input_batch, squeeze_dims=[3]) # reducing the channel dimension.
        input_batch = tf.one_hot(input_batch, depth=n_classes)
    return input_batch

I'm using the tensorflow-deeplab-resnet model which transfers the Resnet model implemented in Caffe to tensorflow using caffe-tensorflow.

Answer 1

According to the documentation, tf.nn.softmax_cross_entropy_with_logits must be called with valid probability distributions on labels, or otherwise the computation will be incorrect, and using tf.nn.sparse_softmax_cross_entropy_with_logits (which might be more convenient in your case) with negative labels will either cause an error or return NaN values. I wouldn't rely on it to have some labels ignored.

What I would do is to replace the logits for the ignored class with infinity in those pixels where the correct class is the ignored one, so they will contribute nothing to the loss:

ignore_label = ...
# Make zeros everywhere except for the ignored label
input_batch_ignored = tf.concat(input_batch.ndims - 1,
    [tf.zeros_like(input_batch[:, :, :, :ignore_label]),
     tf.expand_dims(input_batch[:, :, :, ignore_label], -1),
     tf.zeros_like(input_batch[:, :, :, ignore_label + 1:])])
# Make corresponding logits "infinity" (a big enough number)
predictions_fix = tf.select(input_batch_ignored > 0,
    1e30 * tf.ones_like(predictions), predictions)
# Compute loss with fixed logits
loss = tf.nn.softmax_cross_entropy_with_logits(prediction, gt)

The only problem with this is that you are considering that pixels of the ignored class are always predicted correctly, which means that the loss for images containing a lot of those pixels will be artificially smaller. Depending on the case this may or may not be significant, but if you want to be really accurate, you would have to weight the loss of each image according to the number of not ignored pixels, instead of just taking the mean.

# Count relevant pixels on each image
input_batch_relevant = 1 - input_batch_ignored
input_batch_weight = tf.reduce_sum(input_batch_relevant, [1, 2, 3])
# Compute relative weights
input_batch_weight = input_batch_weight / tf.reduce_sum(input_batch_weight)
# Compute reduced loss according to weights
reduced_loss = tf.reduce_sum(loss * input_batch_weight)

Answer 2

One of the solutions is to write a custom loss function that ignores predictions for labels of your choosing.

I can't provide exact code for semantic segmentation but this is how I tackled the issue when working on a simple named entity recognition problem where I wanted to ignore predictions for specific tokens (like [PAD], [SEP] etc.) when computing the loss during model training.

tokens = ['My', 'name', 'is', 'Robert', '[PAD]']
classes = tf.constant([[0, 0, 0, 1, -1]])
predictions = tf.constant([[-0.5, -0.6, -0.9, 1.2, 0.1]])

This is a binary classification, so BinaryCrossentropy loss can be used:

tf.keras.losses.BinaryCrossentropy(from_logits=True)(classes, predictions)

>>> <tf.Tensor: shape=(), dtype=float32, numpy=0.47207958>

However, just using TensorFlow's BinaryCrossentropy would not ignore predictions for elements with label -1. To ignore them, we can use custom loss function:

def custom_loss(batch_labels, batch_predictions):

    loss = tf.constant(0, dtype=tf.float32)
    binary_crossentropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)
    
    batch_size = tf.shape(batch_labels)[0]
    
    for i in range(batch_size):
        labels = batch_labels[i, :]
        predictions = batch_predictions[i, :]
        
        ix_to_eval = labels != -1
        
        predictions = predictions[ix_to_eval]
        labels = labels[ix_to_eval]
        element_loss = binary_crossentropy(labels, predictions)
        loss = loss + element_loss
        
    return loss/tf.cast(batch_size, tf.float32)

Custom loss takes as an input predictions and labels for one batch of datapoints. Here I iterate over the batch and compute BinaryCrossentropy only for tokens in the sentences with label different from -1.

Note the difference:

custom_loss(classes, predictions)

>>> <tf.Tensor: shape=(), dtype=float32, numpy=0.3790003>

Custom loss can be used in a training loop like any other built-in loss function, e.g.:

model.compile(
    optimizer=optimizer,
    loss=custom_loss)

Answer 3

Sorry, I'm new to this but I believe in https://github.com/tensorflow/models/blob/master/research/deeplab/g3doc/faq.md, it is mentioned that new datasets need to be added here. In "segmentation_dataset.py", at each dataset, you get to specify ignore_label. For example,

_PASCAL_VOC_SEG_INFORMATION = DatasetDescriptor(
    splits_to_sizes={
        'train': 1464,
        'trainval': 2913,
        'val': 1449,
    },
    num_classes=21,
    ignore_label=255,
)