One way of incorporating an underlying metric into the distance of probability measures is to use the Wasserstein distance as the loss - cross entropy loss is the KL divergence - not quite a distance but almost - between the prediction probabilities and the (one-hot distribution given by the labels)  A pytorch implementation and a link to Frogner et al’s paper is linked below.
An alternative could be to use the expected square error loss per pixel. EDIT: The notebook doean’t solve the per pixel problem. For that you might use a Rubner style (http://robotics.stanford.edu/~rubner/papers/rubnerIjcv00.pdf 5) approach to have label as a third dimension, or you might just use MSE…
Best regards
Thomas

  
      github.com
  
  
    t-vi/pytorch-tvmisc/blob/master/wasserstein-distance/Pytorch_Wasserstein.ipynb 16
{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "016D036D46DD4C408D48EFA828B4E77C"
   },
   "source": [
    "# Batch Sinkhorn Iteration Wasserstein Distance\n",
    "\n",
    "Thomas Viehmann\n",
    "\n",
    "This notebook implements sinkhorn iteration wasserstein distance layers.\n",
    "\n",
    "## Important note: This is under construction and does not yet work as well as it should."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,


  This file has been truncated. show original