Manos
Manos

Reputation: 11

Multiple Output Layers Neural Network in TensorFlow

I am currently trying to create a Neural Network in TensorFlow, that has two Output Layers. Specifically I want the network's penultimate layer to serve both as the first Output Layer, but at the same time pass its output to the next and final layer of the network (2nd Output Layer). I need this done in order to apply two different custom losses (one for each output layer). The first one will backpropagate its error starting from the penultimate layer all the way back to the 1st layer, while the second will perform the same, this time backpropagating a different error starting from the network's last layer. Does anyone have any idea how can this be achieved?

The only thing I have achieved is applying a loss function to the penultimate layer and backpropagating the error back to the first layer. This is achieved by creating an identical submodel of the original model, with the difference of having the original model's penultimate layer as its final layer. That way I basically train the submodel the normal way and then I copy its trained weights back to the original model.

This is the only way I can apply a loss function to a model's intermediate layer as accessing an intermediate layer's output during loss calculation is not differentiable, so it throws an error. At least that's how I have it figured out.

Let's say I create this model:

'''

def get_mlp_model(input_shape, num_classes):

    # Initialize Model
    model = tf.keras.models.Sequential() # Sequential NN

    model.add(Dense(units = 1024, activation='relu', input_shape=input_shape, name='fc_1'))
    model.add(Dense(units = 1024, activation='relu', name='fc_2'))
    model.add(Dense(units = 1024, activation='relu', name='fc_3'))
    model.add(Dense(units = 512, activation='relu', name='fc_4'))
    model.add(Dense(units = 512, activation='relu', name='fc_5'))
    model.add(Dense(units = 512, activation='relu', name='fc_6'))
    model.add(Dense(units = num_classes, activation=None, name='fc_7'))

    return model

'''

Then I define this custom loss:

'''

def my_custom_loss(model, x, y_true, layer_name):  

    # Get intermediate layer's output
    get_layer_output = K.function([model.layers[0].input][model.get_layer(layer_name).output]) # Layer Output Function
    y_pred = get_layer_output(x)[0] 

    # Calculate loss using y_true and y_pred...

    return loss

'''

This however throws an error. So in order to achieve calculating the loss based on an intermediate layer, I create an identical submodel as I explained, in which said layer serves as its final layer (output layer) and then train it normally:

'''

def my_custom_loss(model, x, y_true):

    # Get output Layer Output Function
    y_pred = model(x)

    # Calculate loss using y_true and y_pred...
    return loss

'''

But now I need to calculate two separate losses, one for the penultimate layer and the other for the final layer, and apply them to the model simultaneously (after each batch), so the sub-model method doesn't really work anymore.

Upvotes: 1

Views: 7461

Answers (1)

gerwin
gerwin

Reputation: 909

You can use the Funcitonal API to do this. In the functional API, every layer is a variable and you can simply define multiple outputs when finalizing the model.

More info here: https://www.tensorflow.org/alpha/guide/keras/functional

Upvotes: 2

Related Questions