I am trying to implement sample- and pixel-dependent dependent loss weighting in tf.Keras (TensorFlow 2.0.0rc0) for a 3-D U-Net with sparse annotation data (Cicek 2016, arxiv:1606.06650).
This is my code:
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, losses, models
# disabling eager execution makes this example work:
# tf.python.framework_ops.disable_eager_execution()
def get_loss_fcn(w):
    def loss_fcn(y_true, y_pred):
        loss = w * losses.mse(y_true, y_pred)
        return loss
    return loss_fcn
data_x = np.random.rand(5, 4, 1)
data_w = np.random.rand(5, 4)
data_y = np.random.rand(5, 4, 1)
x = layers.Input([4, 1])
w = layers.Input([4])
y = layers.Activation('tanh')(x)
model = models.Model(inputs=[x, w], outputs=y)
loss = get_loss_fcn(model.input[1])
# using another loss makes it work, too:
# loss = 'mse'
model.compile(loss=loss)
model.fit((data_x, data_w), data_y)
print('Done.')
This runs fine when disabling eager execution, but one of the points of TensorFlow 2 is to have eager execution by default. What stands between me and that goal is the custom loss function, as you can see (using 'mse' as a loss removes that error, too):
  File "MWE.py", line 30, in <module>
    model.fit((data_x, data_w), data_y)
[...]
tensorflow.python.eager.core._SymbolicException: Inputs to eager execution function cannot be Keras symbolic tensors, but found [<tf.Tensor 'input_2:0' shape=(None, 4) dtype=float32>]
What can I do to make this kind of structure work with eager execution?
One idea that I had was to concatenate w to the output y and separate y_pred into the original y_pred and w in the loss function, but this is a hack I'd like to avoid. It works, though, with changes marked by # HERE:
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, losses, models
# HERE
def loss_fcn(y_true, y_pred):
    w = y_pred[:, :, -1]  # HERE
    y_pred = y_pred[:, :, :-1]  # HERE
    loss = w * losses.mse(y_true, y_pred)
    return loss
data_x = np.random.rand(5, 4, 1)
data_w = np.random.rand(5, 4, 1)  # HERE
data_y = np.random.rand(5, 4, 1)
x = layers.Input([4, 1])
w = layers.Input([4, 1])  # HERE
y = layers.Activation('tanh')(x)
output = layers.Concatenate()([y, w])  # HERE
model = models.Model(inputs=[x, w], outputs=output)  # HERE
loss = loss_fcn  # HERE
model.compile(loss=loss)
model.fit((data_x, data_w), data_y)
print('Done.')
Any other ideas?