U-Net loss function [cross_entropy] quickly cause gradient-vanishing in high order

[windsor718]

Issue

In the segmentation map, we have much less levee pixels, causing imbalanced learning. We may address this issue by class weight, but we are more interested in the geometrical similarity between predicted images and ground truth, which should create a clearer gradient.

Solution

Use distance between two matrices. Each matrix has shape (N, 2), where N is number of samples and 2-dimension contains the coordinates of pixels labelled as 1 (Useful reference). Actual implementation in tensorflow would be like this.

def pairwise_dist (A, B):  
  """
  Computes pairwise distances between each element of A and B.
  Args:
    A,    [m,d] matrix
    B,    [n,d] matrix
  Returns:
    D,    [m,n] matrix of pairwise distances
  """
  with tf.variable_scope('pairwise_dist'):
    # squared norms of each row in A and B
    na = tf.reduce_sum(tf.square(A), 1)
    nb = tf.reduce_sum(tf.square(B), 1)
    
    # na as a row and nb as a co"lumn vectors
    na = tf.reshape(na, [-1, 1])
    nb = tf.reshape(nb, [1, -1])

    # return pairwise euclidead difference matrix
    D = tf.sqrt(tf.maximum(na - 2*tf.matmul(A, B, False, True) + nb, 0.0))
  return D

You could use tf.where to extract the indices (=coordinates) of images.

[windsor718]

Another option is to use the singular value decomposition to measure similarity between two binary 2d arrays. This will be more strict to the shape of geometry, but insensitive to the distance (i.e., if two shapes are identical but their locations are different, the metric still provide perfect score).

# something like this
sum(abs(svd(A)-svd(B)))