Tensorflow 2.0 concrete function structured_input_signature return value

Question

I'm having some difficulty understanding the return types of structured_input_signature when inspecting a tf.ConcreteFunction.

In the google docs https://www.tensorflow.org/guide/concrete_function#using_a_concrete_function a tuple is returned. For example

@tf.function
def power(a,b):
  print('Tracing "power"\n')
  return a**b

float_power = power.get_concrete_function(
  a = tf.TensorSpec(shape=[], dtype=tf.float32),
  b = tf.TensorSpec(shape=[], dtype=tf.float32))

print(float_power.structured_input_signature)
print(float_power.structured_outputs)

prints

Tracing "power"

((TensorSpec(shape=(), dtype=tf.float32, name='a'), TensorSpec(shape=(), dtype=tf.float32, name='b')), {})
Tensor("Identity:0", shape=(), dtype=float32)

However, when the module is saved and loaded, the output of is slightly different:

float_power_mod = tf.Module()
float_power_mod.float_power = float_power
tf.saved_model.save(float_power_mod, './float_power_mod')

mod_4 = tf.saved_model.load('./float_power_mod')
float_power_func = mod_4.signatures['serving_default']
print(float_power_func.structured_input_signature)

prints

((),
 {'a': TensorSpec(shape=(), dtype=tf.float32, name='a'),
  'b': TensorSpec(shape=(), dtype=tf.float32, name='b')})

What's the logic behind populating the tuple vs the dict in the return tuple of structured_input_signature?

Piston

Answer 1

Short answer

dict allows us to pass keyworded arguments to a function, so that we can tag the real-valued input tensors to the corresponding placeholders accepted by TF.

result = float_power_func(a=tf.constant(2.), b=tf.constant(3.))

Long answer

In order to save a TF model, first, we need to serialise tensors. Under the exported directory you can find a .pb file, that is the protobuf used to serialise your entire model. By model, I mean a collection of tensors and the relationships of these tensors, all of them are captured in protobuf. While TF has provided functions for serialisation and taking your code for example

from tensorflow.python.saved_model import nested_structure_coder

coder = nested_structure_coder.StructureCoder()
signature_proto = coder.encode_structure(float_power.structured_input_signature)
print(signature_proto)

prints

tuple_value {
  values {
    tuple_value {
      values {
        tensor_spec_value {
          name: "a"
          shape {
          }
          dtype: DT_FLOAT
        }
      }
      values {
        tensor_spec_value {
          name: "b"
          shape {
          }
          dtype: DT_FLOAT
        }
      }
    }
  }
  values {
    dict_value {
    }
  }
}

However, the above serialised structure does not satisfy the needs. We cannot assign an input to a key because the return is a tuple.

((TensorSpec(shape=(), dtype=tf.float32, name='a'), TensorSpec(shape=(), dtype=tf.float32, name='b')), {})

As you may realised, the real process of serialising a model is much more complicated, which involves adding new tags and signatures for serving, in-replica and cross-replia context for distribute strategy, among so many other. Regradless of all the complexity, the core is the same, get signatures and serialise them, the code is sourced here

signatures = signature_serialization.canonicalize_signatures(signatures)

The signatures get repacked and the input tensors are moved inside of dict_value as key-value pairs

value {
    canonicalized_input_signature {
      tuple_value {
        values {
          tuple_value {
          }
        }
        values {
          dict_value {
            fields {
              key: "a"
              value {
                tensor_spec_value {
                  name: "a"
                  shape {
                  }
                  dtype: DT_FLOAT
                }
              }
            }
            fields {
              key: "b"
              value {
                tensor_spec_value {
                  name: "b"
                  shape {
                  }
                  dtype: DT_FLOAT
                }
              }
            }
          }
        }
      }
    }

and decoding it you will get

((),
 {'a': TensorSpec(shape=(), dtype=tf.float32, name='a'),
  'b': TensorSpec(shape=(), dtype=tf.float32, name='b')})

Causality