Custom loss function with multiple outputs in tensorflow

Question

my network has two outputs and single input. I am trying to write a custom loss function

$$ Loss = Loss_1(y^{true}_1, y^{pred}_1) + Loss_2(y^{true}_2, y^{pred}_2) $$ I was able to write a custom loss function for a single output. But for multiple output, I am struck. Below I wrote a mwe I tried

def model(input_shape=4, output_shape=3, lr=0.0001):
    """
    single input and multi-output
    loss = custom_loss(out_1_true, out_1_pred)+mse(out_2_true, out_2_pred))
    """
    input_layer = Input(input_shape)
    layer_1 = Dense(units=32, activation='elu', kernel_initializer='he_uniform')(input_layer)
    #outputs
    y_1 = Dense(units=output_shape, activation='softmax', kernel_initializer='he_uniform')(layer_1)
    y_2 = Dense(units=1, activation='linear', kernel_initializer='he_uniform')(layer_1)
    
    def custom_loss(y_true, y_pred):

        # both true value of out_1, out_2 are encoded in y_true
        y_true_1     = y_true[:, :1+output_shape]
        y_true_2     = y_true[:, 1+output_shape:]

        #(this part is wrong...I dont know how)
        y_pred_1, y_pred_2 =  y_pred[:, :1+output_shape], y_pred[:, 1+output_shape:]

        #custorm loss for y_pred_1
        entropy_loss = -y_pred_1 * K.log(y_true_1 + 1e-10)

        #mse for y_pred_2
        mse          = -K.mean(K.square(y_pred_2 - y_true_2))

        #net loss
        loss = entropy_loss + C * mse
        return loss

    Network_model = Model(inputs = input_layer, outputs = [y_1, y_2])
    Network_model.compile(loss = custom_loss, optimizer=RMSprop(lr=lr))

    return Network_model

I think, the main issue lies in spiting the y_pred tensor.

PS: For the purpose of mwe, I have use normal cross entropy loss and mse loss functions in the above code. But however, I have a different cost functions.

Answer 1

import tensorflow as tf
from tensorflow.keras.layers import Input, Dense
from tensorflow.keras import Model
from tensorflow.keras.optimizers import RMSprop
import tensorflow.keras.backend as K
import numpy as np

def model(input_shape=4, output_shape=3, lr=0.0001):
    """
    single input and multi-output
    loss = custom_loss(out_1_true, out_1_pred)+mse(out_2_true, out_2_pred))
    """
    input_layer = Input(input_shape)
    layer_1 = Dense(units=32, activation='elu', kernel_initializer='he_uniform')(input_layer)

    # STEP-1:   name your outputs
    y_1 = Dense(units=output_shape, activation='softmax', kernel_initializer='he_uniform', name='actions')(layer_1)
    y_2 = Dense(units=1, activation='linear', kernel_initializer='he_uniform', name='values')(layer_1)
    
    # STEP-2:    Define the custom custom_loss. Note that the args should be y_true, y_pred
    def custom_loss(y_true, y_pred):
        entropy_loss = -y_pred * K.log(y_true + 1e-10)
        return entropy_loss

    # STEP-3:   Define the dicts for loss and lossweights with keys as the name of the output layers
    LossFunc    =     {'actions':custom_loss, 'values':'mse'}
    lossWeights =     {'actions':0.5, 'values':0.5}


    Network_model = Model(inputs = input_layer, outputs = [y_1, y_2])

    # STEP-4:   complie using LossFunc and lossWeights dicts
    Network_model.compile(optimizer=RMSprop(lr=lr), loss=LossFunc, loss_weights=lossWeights, metrics=["accuracy"])

    return Network_model

#training example:

#model
Network_model = model(4,2)
# input
S = np.reshape([1.,2.,3.,4.], (1, -1))
#true outputs
action, value = np.reshape([0.1], (1, -1)), np.reshape([0.1, 0.9] , (1, -1))
Y = [action, value]


Network_model.fit(S, Y)


1/1 [==============================] - 0s 2ms/step - loss: 6.7340 - actions_loss: 1.1512 - values_loss: 12.3169 - actions_accuracy: 0.0000e+00 - values_accuracy: 1.0000

Answer 2

Multi input/ Multi output Model and learning in tensorflow

import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Concatenate, Add
from tensorflow.keras import Model
from tensorflow.keras.optimizers import RMSprop
import tensorflow.keras.backend as K
import numpy as np

def model(input_shapes=(3,2), lr=0.0001):
    """
    single input and multi-output
    loss = custom_loss(out_1_true, out_1_pred)+mse(out_2_true, out_2_pred))
    """
    input_x = Input(input_shapes[0], name="state")
    input_u = Input(input_shapes[1], name="action")
    input_xu = Concatenate()([input_x, input_u])

    #phi
    layer_1 = Dense(units=32, activation='elu', kernel_initializer='he_uniform')(input_x)
    y = Dense(units=32, activation='elu', kernel_initializer='he_uniform')(layer_1)

    #psi
    layer_1 = Dense(units=32, activation='elu', kernel_initializer='he_uniform')(input_xu)
    w = Dense(units=32, activation='elu', kernel_initializer='he_uniform')(layer_1)

    # Linear system
    Ay = Dense(units=32, activation=None)(y)
    Bw = Dense(units=32, activation=None)(w)
    y_new = Add()([Ay, Bw])

    # phi inv
    layer_1 = Dense(units=32, activation='elu', kernel_initializer='he_uniform', )(y_new)
    x_out = Dense(units=input_shapes[0], activation='elu', kernel_initializer='he_uniform', name = 'x_out')(layer_1)

    # psi inv
    layer_1 = Dense(units=32, activation='elu', kernel_initializer='he_uniform')(w)
    u_out = Dense(units=input_shapes[1], activation='elu', kernel_initializer='he_uniform', name = 'u_out')(layer_1)

    # # STEP-1:   name your outputs
    # y_1 = Dense(units=output_shapes[0], activation='softmax', kernel_initializer='he_uniform', name='actions')(layer_1)
    # y_2 = Dense(units=1, activation='linear', kernel_initializer='he_uniform', name='values')(layer_1)
    
    # # STEP-2:    Define the custom custom_loss. Note that the args should be y_true, y_pred
    # def custom_loss(y_true, y_pred):
    #     entropy_loss = -y_pred * K.log(y_true + 1e-10)
    #     return entropy_loss

    # STEP-3:   Define the dicts for loss and lossweights with keys as the name of the output layers
    LossFunc    =     {'x_out':'mse', 'u_out':'mse'}
    lossWeights =     {'x_out':0.5, 'u_out':0.5}


    Network_model = Model(inputs=[input_x, input_u], outputs = [x_out, u_out])

    # STEP-4:   complie using LossFunc and lossWeights dicts
    Network_model.compile(optimizer=RMSprop(lr=lr), loss=LossFunc, loss_weights=lossWeights, metrics=["accuracy"])

    return Network_model

#training example:

#model
Network_model = model()
# input
S_1 = np.reshape([1.,2.,3.], (1, -1))
S_2 = np.reshape([1.,2.], (1, -1))
S = [S_1, S_2]
#true outputs
action, value = np.reshape([0.1], (1, -1)), np.reshape([0.1, 0.9] , (1, -1))
Y = [S_1, S_2]


Network_model.fit(S, S)