<p>To answer my own question:</p>

<p>The authors of the <a href="https://arxiv.org/pdf/1505.04597.pdf" rel="nofollow noreferrer">U-Net paper</a> used a pre-computed weight-map to handle imbalanced classes.</p>

<p>The Institute for Anstronomy of ETH Zurich provided a <a href="https://tf-unet.readthedocs.io/en/latest/index.html" rel="nofollow noreferrer">Tensorflow-based U-Net package</a> which contains a weighted version of the Softmax function (not sparse but they flatten their labels and logits first):</p>

<pre><code>class_weights = tf.constant(np.array(class_weights, dtype=np.float32))
weight_map = tf.multiply(flat_labels, class_weights)
weight_map = tf.reduce_sum(weight_map, axis=1)
loss_map = tf.nn.softmax_cross_entropy_with_logits_v2(logits=flat_logits, labels=flat_labels)
weighted_loss = tf.multiply(loss_map, weight_map)    
loss = tf.reduce_mean(weighted_loss)
</code></pre>
Tensorflow: Weighted sparse_softmax_cross_entropy for inbalanced classes across a single image
<p>I'm working on a binary semantic segmentation task where the distribution of one class is very smalls across any input image, hence there are only a few pixels which are labeled. When using <a href="https://www.tensorflow.org/api_docs/python/tf/losses/sparse_softmax_cross_entropy" rel="nofollow noreferrer">sparse_softmax_cross_entropy</a>
the over all error is easily decreased when ignoring this class. Now, I'm looking for a way to weight the classes by a coefficient which penalizes missclassifications for the specific class higher compared to the other class.</p>
<p>The doc of the loss function states:</p>
<blockquote>
<p><strong>weights</strong> acts as a coefficient for the loss. If a scalar is provided, then the loss is simply scaled by the given value. If <strong>weights</strong> is a tensor of shape [batch_size], then the loss <strong>weights</strong> apply to each corresponding sample.</p>
</blockquote>
<p>If I understand this correctly, it says that specific sample in a batch get weighted differently compared to others. But this is actually not what I'm looking for. Does anyone know how to implement a weighted version of this loss function where the weights scale the importance of a specific class rather than samples?</p>