Sum of two losses in Keras (Perceptual and MSE)

Question

I want to add perceptual loss in my objective function to the MSE loss. I wrote below code for this:

def custom_objective(y_true, y_pred):

   tosub = K.constant([103.939, 116.779, 123.68])
   y1 = vgg_model(y_pred * 255. - tosub)
   y2 = vgg_model(y_true * 255. - tosub)
   loss2 = K.mean(K.square(y2 - y1), axis=-1)
   loss1 = K.mean(K.square(y_pred - y_true), axis=-1)
   loss = loss1 + loss2

return loss

the problem is that shape of loss1 is something like (BatchSize, 224, 224), but the shape of loss2 is (BatchSize, 7, 7), so it gives me error about incompatible shapes which is right. I want to know how could I add this two properly? should I unravel first? and how?

Answer 1

The loss function should always return a scalar (per sample in the batch or over the whole batch), since we want to minimize it (i.e. you can't minimize a vector, unless you define what you mean by "minimizing a vector"). Therefore, one simple way to reduce this to a scalar is to take the average across all the axes, except the batch axis which is averaged over internally:

loss2 = K.mean(K.square(y2 - y1), axis=[1,2,3])
loss1 = K.mean(K.square(y_pred - y_true), axis=[1,2,3])
loss = loss1 + loss2

Update: Let me clarify that it is OK if the loss function returns a vector or even an n-D array (actually the loss function above returns a vector of length batch_size), but keep in mind that at the end Keras takes the average of returned values and that's the real value of loss (which would be minimized).