How to use the tensors inside a Keras custom loss function?

Viewed 783

I need to train a model with a custom loss function, which shall also update some external function right after the prediction, like this:

def loss_fct(y_true, y_pred):
    global feeder
 
    # Change values of feeder given y_pred
    for value in y_pred:
        feeder.do_something(value)
    
    return K.mean(y_true - y_pred, axis=-1)

However this doesn't work, as TF cannot iterate through tensors in AutoGraph:

OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

My model looks like this

model = Sequential()
model.add(Input(shape=(DIM, )))
model.add(Dense(DIM, activation=None))
model.add(Dense(16, activation=None))
model.add(Dense(4, activation="softmax"))
model.compile(optimizer="adam", loss=loss_fct)
model.summary()

And it is trained like this:

model.fit(x=feeder.feed,
    epochs=18,
    verbose=1,
    callbacks=None,
)

Where feeder.feed is a generator yielding 2 NumPy arrays.


1 Answers

After a lot of research, I have came across this answer. It seems that nothing is wrong with the approach, but it's rather a Tensorflow >= 2.2.0 bug, where Eager Execution is enabled by default.

Finally, to solve this problem use model.compile(..., run_eagerly=True) and iteration and access to tensors during training will be available.

Related