Custom loss function with multiple outputs in tensorflow

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.

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)

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


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)

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