title: Weird Nan loss for custom Keras loss
body: <p>
  I'm trying to implement a custom loss in Keras but can't get it to work.</p>
  
  <p>I have implemented it in numpy and with keras.backend:</p>
  
  <pre class="lang-py prettyprint-override"><code>def log_rmse_np(y_true, y_pred):
      d_i = np.log(y_pred) -  np.log(y_true)
      loss1 = (np.sum(np.square(d_i))/np.size(d_i))
      loss2 = ((np.square(np.sum(d_i)))/(2 * np.square(np.size(d_i))))
      loss = loss1 - loss2
      print('np_loss =  %s - %s = %s'%(loss1, loss2, loss))
      return loss
  
  def log_rmse(y_true, y_pred):
      d_i = (K.log(y_pred) -  K.log(y_true))
      loss1 = K.mean(K.square(d_i))
      loss2 = K.square(K.sum(K.flatten(d_i),axis=-1))/(K.cast_to_floatx(2) * K.square(K.cast_to_floatx(K.int_shape(K.flatten(d_i))[0])))
      loss = loss1 - loss2
      return loss
  </code></pre>
  
  <p>When I test and compare the losses with the following function everything seems to work just fine.</p>
  
  <pre class="lang-py prettyprint-override"><code>def check_loss(_shape):
      if _shape == '2d':
          shape = (6, 7)
      elif _shape == '3d':
          shape = (5, 6, 7)
      elif _shape == '4d':
          shape = (8, 5, 6, 7)
      elif _shape == '5d':
          shape = (9, 8, 5, 6, 7)
  
      y_a = np.random.random(shape)
      y_b = np.random.random(shape)
  
      out1 = K.eval(log_rmse(K.variable(y_a), K.variable(y_b)))
      out2 = log_rmse_np(y_a, y_b)
  
      print('shapes:', str(out1.shape), str(out2.shape))
      print('types: ', type(out1), type(out2))
      print('log_rmse:    ', np.linalg.norm(out1))
      print('log_rmse_np: ', np.linalg.norm(out2))
      print('difference:  ', np.linalg.norm(out1-out2))
      assert out1.shape == out2.shape
      #assert out1.shape == shape[-1]
  
  def test_loss():
      shape_list = ['2d', '3d', '4d', '5d']
      for _shape in shape_list:
          check_loss(_shape)
          print ('======================')
  
  test_loss()
  </code></pre>
  
  <p>The above code prints:</p>
  
  <pre class="lang-py prettyprint-override"><code>np_loss =  1.34490449177 - 0.000229461787517 = 1.34467502998
  shapes: () ()
  types:  &lt;class 'numpy.float32'&gt; &lt;class 'numpy.float64'&gt;
  log_rmse:     1.34468
  log_rmse_np:  1.34467502998
  difference:   3.41081509703e-08
  ======================
  np_loss =  1.68258448859 - 7.67580654591e-05 = 1.68250773052
  shapes: () ()
  types:  &lt;class 'numpy.float32'&gt; &lt;class 'numpy.float64'&gt;
  log_rmse:     1.68251
  log_rmse_np:  1.68250773052
  difference:   1.42057615005e-07
  ======================
  np_loss =  1.99736933814 - 0.00386228512295 = 1.99350705302
  shapes: () ()
  types:  &lt;class 'numpy.float32'&gt; &lt;class 'numpy.float64'&gt;
  log_rmse:     1.99351
  log_rmse_np:  1.99350705302
  difference:   2.53924863358e-08
  ======================
  np_loss =  1.95178217182 - 1.60006871892e-05 = 1.95176617114
  shapes: () ()
  types:  &lt;class 'numpy.float32'&gt; &lt;class 'numpy.float64'&gt;
  log_rmse:     1.95177
  log_rmse_np:  1.95176617114
  difference:   3.78277884572e-08
  ======================
  </code></pre>
  
  <p>I never get an exception when I compile and fit my model with this loss and when I run the model with the 'adam'-loss everything works fine.
  However with this loss keras keeps showing a nan-loss:</p>
  
  <pre class="lang-py prettyprint-override"><code>Epoch 1/10000
   17/256 [&gt;.............................] - ETA: 124s - loss: nan
  </code></pre>
  
  <p>Kind of stuck here... Am I doing something wrong?</p>
  
  <p>Using Tensorflow 1.4 on Ubuntu 16.04</p>
  
  <p><strong>Update:</strong></p>
  
  <p>After a suggestion by Marcin Możejko I updated the code but unfortunately the training loss is still Nan:</p>
  
  <pre class="lang-py prettyprint-override"><code>def get_log_rmse(normalization_constant):
      def log_rmse(y_true, y_pred):
          d_i = (K.log(y_pred) -  K.log(y_true))
          loss1 = K.mean(K.square(d_i))
          loss2 = K.square(K.sum(K.flatten(d_i),axis=-1))/K.cast_to_floatx(2 * normalization_constant ** 2)
          loss = loss1 - loss2
          return loss
      return log_rmse
  </code></pre>
  
  <p>Then the model is compiled via:</p>
  
  <pre class="lang-py prettyprint-override"><code>model.compile(optimizer='adam', loss=get_log_rmse(batch_size))
  </code></pre>
  
  <p><strong>Update 2:</strong></p>
  
  <p>The model summary looks like this:</p>
  
  <pre class="lang-py prettyprint-override"><code>Layer (type)                 Output Shape              Param #   
  =================================================================
  input_2 (InputLayer)         (None, 160, 256, 3)       0         
  _________________________________________________________________
  block1_conv1 (Conv2D)        (None, 160, 256, 64)      1792      
  _________________________________________________________________
  block1_conv2 (Conv2D)        (None, 160, 256, 64)      36928     
  _________________________________________________________________
  block1_pool (MaxPooling2D)   (None, 80, 128, 64)       0         
  _________________________________________________________________
  block2_conv1 (Conv2D)        (None, 80, 128, 128)      73856     
  _________________________________________________________________
  block2_conv2 (Conv2D)        (None, 80, 128, 128)      147584    
  _________________________________________________________________
  block2_pool (MaxPooling2D)   (None, 40, 64, 128)       0         
  _________________________________________________________________
  block3_conv1 (Conv2D)        (None, 40, 64, 256)       295168    
  _________________________________________________________________
  block3_conv2 (Conv2D)        (None, 40, 64, 256)       590080    
  _________________________________________________________________
  block3_conv3 (Conv2D)        (None, 40, 64, 256)       590080    
  _________________________________________________________________
  block3_conv4 (Conv2D)        (None, 40, 64, 256)       590080    
  _________________________________________________________________
  block3_pool (MaxPooling2D)   (None, 20, 32, 256)       0         
  _________________________________________________________________
  block4_conv1 (Conv2D)        (None, 20, 32, 512)       1180160   
  _________________________________________________________________
  block4_conv2 (Conv2D)        (None, 20, 32, 512)       2359808   
  _________________________________________________________________
  block4_conv3 (Conv2D)        (None, 20, 32, 512)       2359808   
  _________________________________________________________________
  block4_conv4 (Conv2D)        (None, 20, 32, 512)       2359808   
  _________________________________________________________________
  block4_pool (MaxPooling2D)   (None, 10, 16, 512)       0         
  _________________________________________________________________
  conv2d_transpose_5 (Conv2DTr (None, 10, 16, 128)       1048704   
  _________________________________________________________________
  up_sampling2d_5 (UpSampling2 (None, 20, 32, 128)       0         
  _________________________________________________________________
  conv2d_transpose_6 (Conv2DTr (None, 20, 32, 64)        131136    
  _________________________________________________________________
  up_sampling2d_6 (UpSampling2 (None, 40, 64, 64)        0         
  _________________________________________________________________
  conv2d_transpose_7 (Conv2DTr (None, 40, 64, 32)        32800     
  _________________________________________________________________
  up_sampling2d_7 (UpSampling2 (None, 80, 128, 32)       0         
  _________________________________________________________________
  conv2d_transpose_8 (Conv2DTr (None, 80, 128, 16)       8208      
  _________________________________________________________________
  up_sampling2d_8 (UpSampling2 (None, 160, 256, 16)      0         
  _________________________________________________________________
  conv2d_2 (Conv2D)            (None, 160, 256, 1)       401       
  =================================================================
  Total params: 11,806,401
  Trainable params: 11,806,401
  Non-trainable params: 0
  </code></pre>
  
  <p><strong>Update 3:</strong></p>
  
  <p>Sample y_true:</p>
  
  <p><a href="https://i.stack.imgur.com/abHvP.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/abHvP.png" alt="Sample of an input image:"></a></p>