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: