tensorflow2.0 - Tensorflow 2.0: tf.feature_columns are triggering tf.function retracing warning

Question

I'm using tensorflow 2.4.0 to train a dnn but I've been getting the following waring after saving and then loading the model:

WARNING:tensorflow:11 out of the last 11 calls to <function recreate_function.<locals>.restored_function_body at 0x7f14b49933b0> triggered tf.function retracing. 
Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. 
For (1), please define your @tf.function outside of the loop. 
For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. 
For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for more details.

The surprising thing is that I don't have any custom functions in my code. After debugging by including only one input at a time, I found that these few lines are triggering the warning.

categorical_column = tf.feature_column.categorical_column_with_vocabulary_list(col_name, ['A', 'B', 'C', 'D'])
feature_columns[col_name] = tf.feature_column.indicator_column(categorical_column=categorical_column)

But there is no way to add the @tf.function decorator to this code, since these are not custom functions but tf.feature_columns.

Is this the expected behaviour? Is there a way to remove the warnings?

question from:https://stackoverflow.com/questions/66064840/tensorflow-2-0-tf-feature-columns-are-triggering-tf-function-retracing-warning

Answer 1

Waitting for answers