Pixel-wise weighted loss function in Keras - TensorFlow 2.0

Question

I'm trying to write a pixel-wise weighted loss function for my model written in Keras but in TensorFlow 2.0 it seems that it is not possible anymore, i.e. it is not possible to have a loss function with other inputs than y_true and y_pred

I used to write it as follows:

from tensorflow.keras.layers import Input, Conv2D
from tensorflow.keras.optimizers import Adam
from tensorflow.keras import backend as K

def my_keras_model():
  input = Input((256,256,1), name='input')
  weight = Input((256,256,1), name='weights')
  c1 = Conv2D(16, (3, 3), activation='relu', kernel_initializer='glorot_uniform', padding='same')(input)
  outputs = Conv2D(1, (1, 1), activation='sigmoid')(c1)

  model=Model(input=[input,weight], output=outputs)
  model.compile(optimizer=Adam(learning_rate=0.001, name='adam'), loss=my_weighted_loss(weight))
  return model

def my_weighted_loss(weight):    
   def loss(y_true, y_pred):
     return K.mean(weight * K.binary_crossentropy(y_true, y_pred), axis=-1)
   return loss

Any idea of how to do it in TF 2?

One solution is to disable the eager mode: `tf.compat.v1.disable_eager_execution()`. — today, Apr 15 '20 at 14:04

score 1 · Answer 1 · answered Apr 15 '20 at 10:05

1

One "hacky" way of implementing this would be adding the original input to the output, and writing your own loss function. This way you can do

weight = y_true[...,0]
y_true = y_true[...,1:]

I would also love to hear a better answer :)

answered Apr 15 '20 at 10:05

rmeertens

4,383
3
17
42

score 0 · Answer 2 · answered Dec 28 '20 at 16:15

Actually it is possible to implement weight maps and do it computation inside the model. Since it is binary cross_entropy

model=Model(inputs=[image,weight,mask], outputs=outputs)

Define your model in such a way that incase if your using tf dataset.

output_types=((tf.float32,tf.float32,tf.float32),tf.float32)
output_shapes=(([1024,1024,1],[1024,1024,1],[1024,1024,1]),[1024,1024,1])

Now compute the loss function inside the model

bce = y_true * K.log(y_pred+epsilon) + (1-y_true) * K.log(1-y_pred+epsilon) #you have values of y_true also

Here model output would be this computed loss. Incase if you need a computation out of your model. Just use a Lambda layer for the weights.

weights_out=layers.Lambda(lambda x:x)(weights)

and then output this layer also from your model. So model would have 2 outputs to compute the loss in the form of a tuple and this way also pixelwise weighted loss can be calculated.

 model=Model(inputs=[image,weights,mask], outputs=[outputs,weighted_out])

Pixel-wise weighted loss function in Keras - TensorFlow 2.0

2 Answers2