Using different loss functions for different outputs simultaneously Keras?

Question

I'm trying to make a network that outputs a depth map, and semantic segmentation data separately.

In order to train the network, I'd like to use categorical cross entropy for the segmentation branch, and mean squared error for the branch that outputs the depth map.

I couldn't find any info on implementing the two loss functions for each branches in the Keras documentation for the Functional API.

Is it possible for me to use these loss functions simultaneously during training, or would it be better for me to train the different branches separately?

Answer 1

From the documentation of Model.compile:

loss: String (name of objective function) or objective function. See losses. If the model has multiple outputs, you can use a different loss on each output by passing a dictionary or a list of losses. The loss value that will be minimized by the model will then be the sum of all individual losses.

If your output is named, you can use a dictionary mapping the names to the corresponding losses:

x = Input((10,))
out1 = Dense(10, activation='softmax', name='segmentation')(x)
out2 = Dense(10, name='depth')(x)
model = Model(x, [out1, out2])
model.compile(loss={'segmentation': 'categorical_crossentropy', 'depth': 'mse'},
              optimizer='adam')

Otherwise, use a list of losses (in the same order as the corresponding model outputs).

x = Input((10,))
out1 = Dense(10, activation='softmax')(x)
out2 = Dense(10)(x)
model = Model(x, [out1, out2])
model.compile(loss=['categorical_crossentropy', 'mse'], optimizer='adam')