loss_pspnet.py

raw_output = net.layers['conv6']
 
# Step1 : 
raw_prediction = tf.reshape(raw_output, [-1, args.num_classes])
label_proc = prepare_label(label_batch, tf.stack(raw_output.get_shape()[1:3]), num_classes=args.num_classes, one_hot=False) # [batch_size, h, w]
raw_gt = tf.reshape(label_proc, [-1,])
indices = tf.squeeze(tf.where(tf.less_equal(raw_gt, args.num_classes - 1)), 1)
 
# Ste2 : 
gt = tf.cast(tf.gather(raw_gt, indices), tf.int32)    
prediction = tf.gather(raw_prediction, indices)
 
# Step3 : Pixel-wise softmax loss.
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction, labels=gt)