title: my simple loss function cause NAN
body: <p>I have write a customer loss for myself, but after several steps, the loss became nan, my code is </p>
  
  <pre><code>def my_loss(label_batch, logits_batch, alpha=1.3, beta=0.5):
      softmax_logits_batch = tf.nn.softmax(logits_batch, axis=-1)
  
      indices_not_0 = tf.where(tf.not_equal(label_batch, 0))  # not-zero indices
      indices_0 = tf.where(tf.equal(label_batch, 0))  # zero indices
  
      predict_not_0 = tf.gather_nd(softmax_logits_batch, indices_not_0)
      predict_0 = tf.gather_nd(softmax_logits_batch, indices_0)
      avg_p_not_0 = tf.reduce_mean(predict_not_0, axis=0)
      avg_p_0 = tf.reduce_mean(predict_0, axis=0)
  
      euclidean_distance = tf.sqrt(tf.reduce_sum(tf.square(avg_p_0 - avg_p_not_0)))
      max_value = tf.maximum(alpha - euclidean_distance, 0)
      return max_value
  </code></pre>
  
  <p>Some basic ideas behind are:</p>
  
  <ol>
  <li><p>My loss is for semantic segmentation which has only <strong>2 categories</strong>.</p></li>
  <li><p>The shape of label_batch is <strong>(?, H, W)</strong>, all the value in it are 0 or 1. The shape of logits_batch is <strong>(?, H, W, 2)</strong> the value of logits_batch is the <strong>logits of the FCN (without Softmax)</strong>.</p></li>
  <li><p>I want to find <strong>all</strong> the logits values (<strong>predict_0</strong> or <strong>predict_not_0</strong>) whose label value is 0 or 1 respectively by <strong>indices_0</strong> or <strong>indices_not_0</strong>. </p></li>
  <li><p>The shape of both predict_not_0 and predict_0 should be <strong>(?, 2)</strong></p></li>
  <li><p>Calculate the average value for predict_not_0 and predict_0 respectively (which represents the central point coordinates of Euclidean space for category 0 and category 1). The shape of them should be <strong>(2,)</strong></p></li>
  <li><p>calculate the Euclidean distance between two central point coordinates, and it should larger than a certain <strong>alpha</strong> value (alpha = 1.3 for example)</p></li>
  </ol>
  
  <p>Now, the problem is after several steps, the loss value become nan.</p>
  
  <p>the output of the code is (I used a very small learning rate parameter)</p>
  
  <pre><code>Epoch[0],step[1],train batch loss = 2.87282,train acc = 0.486435.
  Epoch[0],step[2],train batch loss = 2.87282,train acc = 0.485756.
  Epoch[0],step[3],train batch loss = 2.87281,train acc = 0.485614.
  Epoch[0],step[4],train batch loss = 2.87282,train acc = 0.485649.
  Epoch[0],step[5],train batch loss = 2.87282,train acc = 0.485185.
  Epoch[0],step[6],train batch loss = 2.87279,train acc = 0.485292.
  Epoch[0],step[7],train batch loss = 2.87281,train acc = 0.485222.
  Epoch[0],step[8],train batch loss = 2.87282,train acc = 0.484989.
  Epoch[0],step[9],train batch loss = 2.87282,train acc = 0.48406.
  Epoch[0],step[10],train batch loss = 2.8728,train acc = 0.483306.
  Epoch[0],step[11],train batch loss = 2.87281,train acc = 0.483426.
  Epoch[0],step[12],train batch loss = 2.8728,train acc = 0.482954.
  Epoch[0],step[13],train batch loss = 2.87281,train acc = 0.482535.
  Epoch[0],step[14],train batch loss = 2.87281,train acc = 0.482225.
  Epoch[0],step[15],train batch loss = 2.87279,train acc = 0.482005.
  Epoch[0],step[16],train batch loss = 2.87281,train acc = 0.48182.
  Epoch[0],step[17],train batch loss = 2.87282,train acc = 0.48169.
  Epoch[0],step[18],train batch loss = 2.8728,train acc = 0.481279.
  Epoch[0],step[19],train batch loss = 2.87281,train acc = 0.480878.
  Epoch[0],step[20],train batch loss = 2.87281,train acc = 0.480607.
  Epoch[0],step[21],train batch loss = 2.87278,train acc = 0.480186.
  Epoch[0],step[22],train batch loss = 2.87281,train acc = 0.479925.
  Epoch[0],step[23],train batch loss = 2.87282,train acc = 0.479617.
  Epoch[0],step[24],train batch loss = 2.87282,train acc = 0.479378.
  Epoch[0],step[25],train batch loss = 2.87281,train acc = 0.479496.
  Epoch[0],step[26],train batch loss = 2.87281,train acc = 0.479354.
  Epoch[0],step[27],train batch loss = 2.87282,train acc = 0.479262.
  Epoch[0],step[28],train batch loss = 2.87282,train acc = 0.479308.
  Epoch[0],step[29],train batch loss = 2.87282,train acc = 0.479182.
  Epoch[0],step[30],train batch loss = 2.22282,train acc = 0.478985.
  Epoch[0],step[31],train batch loss = nan,train acc = 0.494112.
  Epoch[0],step[32],train batch loss = nan,train acc = 0.508811.
  Epoch[0],step[33],train batch loss = nan,train acc = 0.523289.
  Epoch[0],step[34],train batch loss = nan,train acc = 0.536233.
  Epoch[0],step[35],train batch loss = nan,train acc = 0.548851.
  Epoch[0],step[36],train batch loss = nan,train acc = 0.561351.
  Epoch[0],step[37],train batch loss = nan,train acc = 0.573149.
  Epoch[0],step[38],train batch loss = nan,train acc = 0.584382.
  Epoch[0],step[39],train batch loss = nan,train acc = 0.595006.
  Epoch[0],step[40],train batch loss = nan,train acc = 0.605065.
  Epoch[0],step[41],train batch loss = nan,train acc = 0.614475.
  Epoch[0],step[42],train batch loss = nan,train acc = 0.623371.
  Epoch[0],step[43],train batch loss = nan,train acc = 0.632092.
  Epoch[0],step[44],train batch loss = nan,train acc = 0.640199.
  Epoch[0],step[45],train batch loss = nan,train acc = 0.647391.
  </code></pre>
  
  <p><strong>I used exactly the same code before, except the loss function is tf.nn.sparse_softmax_cross_entropy_with_logits() and everything works, so I suppose there is something wrong in my new loss function.</strong></p>
  
  <p>I have a guess, maybe some batch data only have one category's label (only 0 or 1), so one of <strong>predict_not_0 and predict_0</strong> will have no data therefore, but I don't know how to code to validate whether there has data in <strong>predict_not_0 and predict_0</strong></p>
  
  <p>can somebody help me find where the problem is and how can I improve my loss function to avoid nan?</p>