How do I print inside the loss function during training in Keras?

Question

I am trying to create a loss function in Keras (Tensorflow Backend) but I am a little stuck to check the inside of the custom loss function. In fact, the print appears on the console only when I compile the model, after that there is no print. (I am just testing very simple custom function, I will create the true function when I solved this problem). I train the model using the train_on_batch function. How can I solve this problem?

def loss_yolo(self, y_true, y_pred):  
    print('inside loss function')
    loss = K.mean(y_true - y_pred)
    return loss

model.compile(optimizer='sgd', loss=loss_yolo)

print('train on batch')
print(model.train_on_batch(x, y))

And the output is

inside loss function

train on batch

-0.481604

Answer 1

The only thing you can do is not use python's print function, but for example, tensorflow's tf.Print function that is part of the computational graph. The documentation says the operation does nothing but each time it is evaluated it prints a message that you can specify.

You just need to be careful to place it correctly in the graph, something like:

def loss(y_true, y_pred):
    d = y_true - y_pred
    d = tf.Print(d, [d], "Inside loss function")
    return tf.reduce_mean(tf.square(d))

A better option to look inside what is going on internally is to use the tensorflow debugger.

Answer 2

I added the output_stream argument and tried this code in TensorFlow v2.4.1. Worked fine:

def loss_custom(y_true, y_pred):
    d = y_true - y_pred
    tf.print("\n y_true:", type(y_true), output_stream=sys.stdout)
    return tf.reduce_mean(tf.square(d))

Output during training:

Epoch 1/10

 y_true: <class 'tensorflow.python.framework.ops.Tensor'>
 1/72 [..............................] - ETA: 0s - loss: 0.2328 - accuracy: 0.3319
 y_true: <class 'tensorflow.python.framework.ops.Tensor'>
 2/72 [..............................] - ETA: 9s - loss: 0.2087 - accuracy: 0.5250
 y_true: <class 'tensorflow.python.framework.ops.Tensor'>