Hands-On Convolutional Neural Networks with TensorFlow

The loss function

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: