MSE vs Cross Entropy for training with facial landmark (pose) heatmaps

Question

I am trying to reimplement the excellent paper https://github.com/1adrianb/face-alignment-training in tensorflow. I have successfully defined the network and downloaded the LSD3D-W dataset. I am able to train the model however I am running into a serious issue.

Ground Truth and Loss

For training, I generate ground truths by converting the x,y landmark coordinates into gaussian heatmaps where x,y is the mean of the gaussian.
I first trained with MSE loss as given in the original implementation. After some iterations, the loss becomes extremely small but the output is completely white!

loss = tf.losses.mean_squared_error(
    predictions=heatmaps, labels=labels_tensor
)

When I tried with cross entropy, I am getting better results. But they are not sharper

loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=heatmaps, labels=labels_tensor), name= 'cross_entropy_loss')

gtmap means groundtruth map.

Answer 1

I experienced a similar problem while doing pose estimation. Your problem may be due to the huge amount of "background" pixels, where the gaussian heatmap has very small values. So the network simply learns the trivial solution and sets the whole heatmap to zero, which results in a small mse error.

Solution to my problem was:

1) weigh the heatmaps: weigh pixels in the loss function, which are "foreground" more. For example you could count the percentage of "foreground" pixels and the percentage of "background" pixels for each heatmap and weigh foreground pixels with (1-percentage of foreground pixels) and background pixels with (1-percentage of background pixels)

2) Increase the maximum of the gaussians to ~12