Keras custom loss function with three network inputs and model.add_loss

Hi I need some help with a custom loss function in Keras. I am basically building a UNET with a second input which takes on the weight maps like in the original UNET paper.
However I am using this UNET for image synthesis and my loss function is a combination of perceptual loss and pixel loss computed using three inputs (input image, reconstructed image and weight maps). The UNET model is a standard UNET with encoder, decoder and skip connections.

Below is my code for the network and for the loss function:

def synthesis_unet_weights(pretrained_weights=None, input_shape=(SIZE_s, SIZE_s, 3), num_classes=1, is_training=True):
    
    ip        = Input(shape=input_shape)
    weight_ip = Input(shape=input_shape[:2] + (num_classes,))
    
    UNET encoder with the first Conv2D layer taking input ip
#---------------------------------------------------------------------------------------------------------------------------    
    center = Conv2D(1024, (3,3),padding='same', activation='relu', kernel_initializer=initializer)(pool4)
    center = Conv2D(1024, (3,3),padding='same', activation='relu', kernel_initializer=initializer)(center)
#---------------------------------------------------------------------------------------------------------------------------
    UNET decoder with the last layer up1
    
    classify = Conv2D(num_classes, (1,1), activation='sigmoid')(up1)
    
    if is_training:
               
        model=Model(inputs=[ip, weight_ip], outputs=[classify])
        model.add_loss(perceptual_loss_weight(ip,classify,weight_ip))
         
        return model
     
    else:
        
        model = Model(inputs=[ip], outputs=[classify])
        weight_ip=ip
        model.add_loss(perceptual_loss_weight(ip,classify,weight_ip))
        
        opt2 = tf.keras.optimizers.Adam(learning_rate=1e-3,clipnorm=1.0)
        model.compile(optimizer=opt2)
        
        return model        
    return model



def perceptual_loss_weight(input_image , reconstruct_image,  weights):
                       
    input_image       = clip_0_1(input_image)
    reconstruct_image = tf.concat((reconstruct_image,reconstruct_image,reconstruct_image),axis=-1)
    reconstruct_image = clip_0_1(reconstruct_image)
    weights = tf.concat((weights,weights,weights),axis=-1)
    weights = clip_0_1(weights)
            
    h1_list = LossModel(input_image)
    h2_list = LossModel(reconstruct_image)
   
    rc_loss = 0.0

    for h1, h2, weight in zip(h1_list, h2_list, selected_layer_weights):
        
           h1 = K.batch_flatten(h1)
           h2 = K.batch_flatten(h2)
                
           rc_loss = rc_loss + weight * K.sum(K.square(h1 - h2), axis=-1)
           
    
    pixel_loss = K.sum(K.square(K.batch_flatten(weights)*K.batch_flatten(input_image) - K.batch_flatten(weights)*K.batch_flatten(reconstruct_image)),axis=1) 
    return rc_loss+pixel_loss

The weight input is used only for the loss function during training. I managed to train the model (compiling with loss=None) but it is not predicting what it should predict. It looks like the input is just passed through the network (without any modification) directly to the output. The reconstructed output image looks exactly the same as the input image.

Okay, I have found the conceptual mistake, I was feeding the input images for input_image input to the loss function. But actually those should be y_true labels. One possible solution is to feed the network an additional input “ip_labels” which are the y_true required by the perceptual loss @tf.function. Below is the working solution written as a custom loss function with a dummy loss layer:

def MyLoss2(input_image, reconstruct_image, weight_ip):
@tf.function
def perceptual_loss(input_image, reconstruct_image):
input_image       = clip_0_1(input_image)
reconstruct_image = clip_0_1(reconstruct_image)
weights = clip_0_1(weight_ip)
h1_list = LossModel(input_image)
h2_list = LossModel(reconstruct_image)
rc_loss = 0.0
for h1, h2, weight in zip(h1_list, h2_list, selected_layer_weights):
h1 = K.batch_flatten(h1)
h2 = K.batch_flatten(h2)
rc_loss = rc_loss + weight * K.sum(K.square(h1 - h2), axis=-1)
pixel_loss = K.sum(K.square(K.batch_flatten(weights)*K.batch_flatten(input_image) - K.batch_flatten(weights)*K.batch_flatten(reconstruct_image)),axis=1)
return rc_loss + pixel_loss
return perceptual_loss(input_image, reconstruct_image)

The loss layer is implemented as follows:

class DummyLayer(Layer):
def __init__(self, is_training):
super().__init__(is_training)
self.is_training=is_training
def get_config(self):
config = super().get_config()
config.update({
"is_training": self.is_training,
})
return config
def call(self, inputs,is_training):
ip,classify,weight_ip = tf.unstack(inputs,axis=-1)
self.add_loss(MyLoss2(ip, classify, weight_ip))
return inputs

The synthesis UNET has an additional input ip_labels:

def synthesis_unet_weights(pretrained_weights=None, input_shape=(SIZE_s, SIZE_s, 3), num_classes=1, is_training=True):
ip        = Input(shape=input_shape)
ip_labels = Input(shape=input_shape)
weight_ip = Input(shape=input_shape[:2] + (num_classes,))
down1 = Conv2D(64, (3,3),padding='same', activation=LeakyReLU(alpha=0.3), kernel_initializer=initializer)(ip)
UNET encoder with skips
center = Conv2D(1024, (3,3),padding='same', activation='relu', kernel_initializer=initializer)(pool4)
center = Conv2D(1024, (3,3),padding='same', activation='relu', kernel_initializer=initializer)(center)
UNET decoder with the last layer up1
classify = Conv2D(num_classes, (1,1), activation='sigmoid')(up1)
if is_training:
data=tf.stack([ip_labels,(tf.concat((classify,classify,classify),axis=3)),
(tf.concat((weight_ip,weight_ip,weight_ip),axis=3))],axis=-1)
classify = DummyLayer(is_training=True)(data, is_training=True)
inp, classify, weight_inp=tf.unstack(classify,axis=-1)
model=Model(inputs=[ip, ip_labels, weight_ip], outputs=[classify])
opt = tf.keras.optimizers.Adam(learning_rate=1e-3,clipnorm=1.0)
model.compile(optimizer=opt, metrics=['mse','mae'])      
return model
else:
data=tf.stack([ip,(tf.concat((classify,classify,classify),axis=3)), ip],axis=-1)
classify=DummyLayer(is_training=False)(data, is_training=False)
inp, classify, weight_inp=tf.unstack(classify,axis=-1)
model = Model(inputs=[ip], outputs=[classify])
opt = tf.keras.optimizers.Adam(learning_rate=1e-3,clipnorm=1.0)
model.compile(optimizer=opt)
return model
return model

To train :

model=synthesis_unet_weights()
model.fit([input_images, labels, weight_maps], labels)

I know the code with stacking and unstacking tensors is not particularly elegant but it is working.