Пиксельная взвешенная функция потерь в Керасе - TensorFlow 2.0

Question

Я пытаюсь написать побитовую взвешенную функцию потерь для моей модели, написанной на Keras, но в TensorFlow 2.0 кажется, что это больше невозможно, то есть невозможно иметь функцию потерь с другими входами, чем y_true и y_pred

Я писал так:

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

Есть идеи, как это сделать в TF 2?

Answer 1

Один "хакерский" способ реализовать это - добавить исходный ввод к выводу и написать собственную функцию потерь. Таким образом, вы можете сделать

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

Я также хотел бы услышать лучший ответ :)