Keras custom loss function vs. Lambda layer

Viewed 1869

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.

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.

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.

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:

#-----------------------------------------------------
# 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
#-----------------------------------------------------
# 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

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.

1 Answers

Use keepdims=True in the K.sum operations. This is required to maintain the correct shape.

Try the following:

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import *
from tensorflow.keras.models import Model
from tensorflow.keras import backend as K

BINSIZE = 1
XMIN = 0

def weighted_avg(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, keepdims=True)          # normalization of all outputs by batch. shape is 1D array of size batch
    wavg = K.sum(idx*inputs, axis=-1, keepdims=True)/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

def make_model():
  inp = Input(shape=(4,))
  out = Lambda(weighted_avg)(inp)
  model = Model(inp, out)
  model.compile('adam', 'mse')
  return model

model = make_model()
model.summary()

Simple test code:

import numpy as np
X = np.array([
    [1, 1, 1, 1],
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 1, 0],
    [0, 0, 0, 1]
])
model.predict(X)

predict should emit a column vector such as:

array([[2.5],
       [1. ],
       [2. ],
       [3. ],
       [4. ]], dtype=float32)
Related