As mentioned, the loss function for segmentation models will basically be an extension of the classification loss, but working spatially throughout the whole output vector:
# Segmentation problems often uses this "spatial" softmax (Basically we want to classify each pixel) with tf.name_scope("SPATIAL_SOFTMAX"): loss = tf.reduce_mean((tf.nn.sparse_softmax_cross_entropy_with_logits( logits=model_out,labels=tf.squeeze(labels_in, squeeze_dims=[3]),name="spatial_softmax")))
The below image depicts the implementation of Fully Convolutional Networks for Semantic Segmentation:
The below image shows the SegNet architecture: