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?

Answer 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 :)