title: TensorFlow 2: Re-saving a SavedModel?
body: <p>I am trying to load a model saved in <a href="https://www.tensorflow.org/guide/saved_model" rel="nofollow noreferrer">SavedModel</a> format, and then apply some calculations on top of it and re-save the whole pipeline. A minimal code is the following (from <a href="https://www.kaggle.com/camaskew/host-baseline-example" rel="nofollow noreferrer">this Kaggle kernel</a>):</p>
  <pre class="lang-python prettyprint-override"><code># 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'])
  </code></pre>
  <blockquote>
  <p><strong>Question:</strong> Now, I want to save the <code>global_feature_extraction_fn</code>, with some other TF ops to post-process the output, in the same SavedModel format. What is the correct way to do that?</p>
  </blockquote>
  <hr />
  <h2>What I have tried</h2>
  <p>I have tried to follow the TensorFlow documentation of <a href="https://www.tensorflow.org/guide/saved_model#saving_a_custom_model" rel="nofollow noreferrer">saving custom models to SavedModel format</a> and define a <code>tf.Module</code>:</p>
  <pre class="lang-python prettyprint-override"><code>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()
  </code></pre>
  <p>Then I ran it on the test image to build the <code>tf.function</code> and make sure that it works correctly (produces the right output). But when I tried to save it as follows:</p>
  <pre class="lang-python prettyprint-override"><code>tf.saved_model.save(
      delg_module, export_dir='./delg_resaved',
      signatures={
          'serving_default': delg_module.call
      })
  </code></pre>
  <p>the resulting model was incorrect. The original one weighted 90 MB while the one in <code>delg_resaved</code> weights only 800 KB. I also tried doing <code>tf.saved_model.load</code> <strong>inside</strong> the <code>DelgModule.call</code> function, so that the graph creation and variables loading is done entirely inside that <code>tf.function</code>, but the results remained the same.</p>