Use tf1 model in tf2 pipeline with tf.GradientTape(persistent=True)

TensorFlow General Discussion
,
1

I have a tf1 model which I want to use in a tf2 training pipeline with tf.GradientTape(persistent=True).

The intuitive method to do so would be to convert the tf1 model to tf2 and save it as a SavedModel, as explained in tf documentation to migrate from tf1 to tf2 (which I tried).

However, as highlighted by this comment , when saving a SavedModel with a tf version anterior to 2.5, the model cannot be used with tf.GradientTape(persistent=True). Although I am using tf 2.9 for the conversion, the methods used to save the tf1 model are from tf.compat.v1 and hence do not include the fix introduced in tf 2.5.

Code example:

import tensorflow as tf

MODEL_FOLDER_PATH = "" # any model saved with tf1

model = tf.saved_model.load(MODEL_FOLDER_PATH, tags="serve").signatures[
    "serving_default"
] # loading a model which has been saved as explained in tf documentation:
# "Migrating the SavedModel workflow": https://www.tensorflow.org/guide/migrate/saved_model

@tf.function
def not_using_persistent_gradient(model):
    with tf.GradientTape():
        model(tf.ones((1, 64, 256, 3)))

@tf.function
def using_persistent_gradient(model):
    with tf.GradientTape(persistent=True):
        model(tf.ones((1, 64, 256, 3)))

not_using_persistent_gradient(model)   # This works!
using_persistent_gradient(model)  # This raises an error!

The error is:

ValueError: Internal error: Tried to take gradients (or similar) of a variable without handle data:
Tensor("Backward/Predictor/decoder/while:13", shape=(), dtype=resource)

Would anyone know how to make it work?