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.