title: Keras custom loss function vs. Lambda layer
body: <p>I have model that I am able to train using a custom loss function and it works fine. I would like to replace the customized loss function with the standard mean_squared_error by moving some calculations to a Lambda layer instead.</p>
  
  <p>Some details:
  The model ultimately needs to produce a single floating point number. The original model had 60 outputs which I converted to a single number by taking a weighted average. I did this in the loss function to compare to the label, but then also have to do it after inference. I would like to embed doing this weighted average into the last layer of the network itself which would then simplify things.</p>
  
  <p>I would like to head of suggestions to just add a single node Dense layer to the end and let the network figure it out. I tried it and it does not work very well. (I believe the problem being that the weighted average requires a division operation which would need to be mimic'd by a couple of more Dense layers). Regardless, I'm really interested in understanding the Lambda layers so I can add it to my toolbox.</p>
  
  <p>Here is some code that shows the two things I've done. I've minimized it as much as possible. These are snippets from larger scripts, but the parts not shown are identical for those and these are the only differences:</p>
  
  <pre class="lang-py prettyprint-override"><code>#-----------------------------------------------------
  # customLoss
  #-----------------------------------------------------
  # Define custom loss function that compares calcukated phi
  # to true
  def customLoss(y_true, y_pred):
  
      # Calculate weighted sum of prediction
      ones = K.ones_like(y_pred[0,:])       # [1, 1, 1, 1....]   (size Nouts)
      idx  = K.cumsum(ones)                 # [1, 2, 3, 4....]   (size Nouts)
      norm = K.sum(y_pred, axis=1)          # normalization of all outputs by batch. shape is 1D array of size batch
      wavg = K.sum(idx*y_pred, axis=1)/norm # array of size batch with weighted avg. of mean in units of bins
      wavg_cm = wavg*BINSIZE + XMIN         # array of size batch with weighted avg. of mean in physical units
  
      # Calculate loss
      loss_wavg = K.mean(K.square(y_true[:,0] - wavg_cm), axis=-1)
  
      return loss_wavg
  
  #-----------------------------------------------------
  # DefineModel
  #-----------------------------------------------------
  # This is used to define the model. It is only called if no model
  # file is found in the model_checkpoints directory.
  def DefineModel():
  
      # Build model
      inputs = Input(shape=(height, width, 1), name='image_inputs')
      x = Flatten()(inputs)
      x = Dense( int(Nouts*5), activation='linear')(x)
      x = Dense( Nouts, activation='relu')(x)
      model = Model(inputs=inputs, outputs=[x])
  
      # Compile the model and print a summary of it
      opt = Adadelta(clipnorm=1.0)
      model.compile(loss=customLoss, optimizer=opt)
  
      return model
  </code></pre>
  
  <pre class="lang-py prettyprint-override"><code>#-----------------------------------------------------
  # MyWeightedAvg
  #
  # This is used by the final Lambda layer of the network.
  # It defines the function for calculating the weighted
  # average of the inputs from the previous layer.
  #-----------------------------------------------------
  def MyWeightedAvg(inputs):
  
      # Calculate weighted sum of inputs
      ones = K.ones_like(inputs[0,:])       # [1, 1, 1, 1....]   (size Nouts)
      idx  = K.cumsum(ones)                 # [1, 2, 3, 4....]   (size Nouts)
      norm = K.sum(inputs, axis=1)          # normalization of all outputs by batch. shape is 1D array of size batch
      wavg = K.sum(idx*inputs, axis=1)/norm # array of size batch with weighted avg. of mean in units of bins
      wavg_cm = wavg*BINSIZE + XMIN         # array of size batch with weighted avg. of mean in physical units
  
      return wavg_cm
  
  #-----------------------------------------------------
  # DefineModel
  #-----------------------------------------------------
  # This is used to define the model. It is only called if no model
  # file is found in the model_checkpoints directory.
  def DefineModel():
  
      # Build model
      inputs = Input(shape=(height, width, 1), name='image_inputs')
      x = Flatten()(inputs)
      x = Dense( int(Nouts*5), activation='linear')(x)
      x = Dense( Nouts, activation='relu')(x)
      x = Lambda(MyWeightedAvg, output_shape=(1,), name='z_output')(x)
      model = Model(inputs=inputs, outputs=[x])
  
      # Compile the model and print a summary of it
      opt = Adadelta(clipnorm=1.0)
      model.compile(loss='mean_squared_error', optimizer=opt)
  
      return model
  </code></pre>
  
  <p>I expected these to give the same result, but the custom loss function seems to train well and produce loss values that drop fairly steadily over several epochs while the Lamda drops to value of 18.72... and kind of oscillates close to that.</p>