title: Weighted Pixel Wise Categorical Cross Entropy for Semantic Segmentation
body: <p>I have recently started learning about Semantic Segmentation. I am trying to train a UNet for the same. My input is RGB 128x128x3 images. My masks are made up of 4 classes 0, 1, 2, 3 and are One-Hot Encoded with dimension 128x128x4.</p>
  
  <pre><code>def weighted_cce(y_true, y_pred):
          weights = []
          t_inf = tf.convert_to_tensor(1e9, dtype = 'float32')
          t_zero = tf.convert_to_tensor(0, dtype = 'int64')
          for i in range(0, 4):
              l = tf.argmax(y_true, axis = -1) == i
              n = tf.cast(tf.math.count_nonzero(l), 'float32') + K.epsilon()
              weights.append(n)
  
          weights = [batch_size/j for j in weights]
  
          y_pred /= K.sum(y_pred, axis=-1, keepdims=True)
          # clip to prevent NaN's and Inf's
          y_pred = K.clip(y_pred, K.epsilon(), 1 - K.epsilon())
          # calc
          loss = y_true * K.log(y_pred) * weights
          loss = -K.sum(loss, -1)
          return loss
  </code></pre>
  
  <p>This is the loss function that I am using but it classifies every pixel as 2. What am I doing wrong?</p>