TensorFlow 2: Re-saving a SavedModel?

Viewed 322

I am trying to load a model saved in SavedModel format, and then apply some calculations on top of it and re-save the whole pipeline. A minimal code is the following (from this Kaggle kernel):

# Load the model. After that, the `delg_model` will be of the type
# `tensorflow.python.training.tracking.tracking.AutoTrackable`
delg_model = tf.saved_model.load('path/to/saved/model/dir')

# we don't need the whole model, so we prune it. After that, the
# `global_feature_extraction_fn` will be of the type
# `tensorflow.python.eager.wrap_function.WrappedFunction`
delg_input_tensor_names = ['input_image:0', 'input_scales:0']
global_feature_extraction_fn = delg_model.prune(
    delg_input_tensor_names, ['global_descriptors:0'])

Question: Now, I want to save the global_feature_extraction_fn, with some other TF ops to post-process the output, in the same SavedModel format. What is the correct way to do that?


What I have tried

I have tried to follow the TensorFlow documentation of saving custom models to SavedModel format and define a tf.Module:

class DelgModule(tf.Module):
    def __init__(self):
        super().__init__()

    @tf.function(input_signature=[
        tf.TensorSpec(shape=[None, None, 3], name='input_image')
    ])
    def call(self, input_tensor):
        # custom function on top of the model's output
        embedding = tf.nn.l2_normalize(
            global_feature_extraction_fn(
                input_tensor,                         # input_image
                tf.convert_to_tensor([0.7, 1.0, 1.4]) # input_scales
            )[0],
            axis=1, name='l2_normalized_output')
        return output_tensors = {
            'global_descriptor': embedding
        }

delg_module = DelgModule()

Then I ran it on the test image to build the tf.function and make sure that it works correctly (produces the right output). But when I tried to save it as follows:

tf.saved_model.save(
    delg_module, export_dir='./delg_resaved',
    signatures={
        'serving_default': delg_module.call
    })

the resulting model was incorrect. The original one weighted 90 MB while the one in delg_resaved weights only 800 KB. I also tried doing tf.saved_model.load inside the DelgModule.call function, so that the graph creation and variables loading is done entirely inside that tf.function, but the results remained the same.

0 Answers
Related