Search code examples
keraskeras-layer

Keras Lambda Layer for Custom Loss


I am attempting to implement a Lambda layer that will produce a custom loss function. In the layer, I need to be able to compare every element in a batch to every other element in the batch in order to calculate the cost. Ideally, I want code that looks something like this:

for el_1 in zip(y_pred, y_true):
    for el_2 in zip(y_pred, y_true):
        if el_1[1] == el_2[1]:
            # Perform a calculation
        else:
            # Perform a different calculation

When I true this, I get:

TypeError: TensorType does not support iteration.

I am using Keras version 2.0.2 with a Theano version 0.9.0 backend. I understand that I need to use Keras tensor functions in order to do this, but I can't figure out any tensor functions that do what I want.

Also, I am having difficulty understanding precisely what my Lambda function should return. Is it a tensor of the total cost for each sample, or is it just a total cost for the batch?

I have been beating my head against this for days. Any help is deeply appreciated.


Solution

  • A tensor in Keras commonly has at least 2 dimensions, the batch and the neuron/unit/node/... dimension. A dense layer with 128 units trained with a batch size of 64 would therefore yields a tensor with shape (64,128).

    Your LambdaLayer processes tensors as any other layer does, plugging it in after your dense layer from before will give you a tensor with shape (64,128) to process. Processing a tensor works similar to how calculations on numpy arrays works (or any other vector processing library really): you specify one operation to broadcast over all elements in the data structure.

    For example, your custom cost is the difference for each value in the batch, you would implement it like so:

    cost_layer = LambdaLayer(lambda a,b: a - b)
    

    The - operation is broadcasted over a and b and will return a suitable result provided the dimensions match. The takeaway is that you really only can specify one operation for every value. If you want to do more complex tasks, for example computations based on the value you need single operations that take two operations and apply the correct one accordingly, i.e. the switch operation.

    The syntax for K.switch is

    K.switch(condition, then_expression, else_expression)
    

    For example, if you want to subtract both values when a != b but add them when they are equal, you would write:

    import keras.backend as K
    cost_layer = LambdaLayer(lambda a,b: K.switch(a != b, a - b, a + b))