Posted by Daniel Ellis, TensorFlow Engineer
Note: This blog post is aimed at TensorFlow developers who want to learn the details of how graphs and models are stored. If you are new to TensorFlow, you should check out the and .
Models and layers can be loaded from this representation without actually making an instance of the Python class that created it. This is desired in situations where you do not have (or want) a Python interpreter, such as serving at scale or on an edge device, or in situations where the original Python code is not available.
Saved models are represented by two separate, but equally important, parts: the graph, which describes the fixed computation described in code, and the weights, which are the dynamic parameters you trained during training. If you aren't already familiar with this and @tf.function, you should check out the in the we refer to these as polymorphic functions, as they are Python callables that can take a variety argument signatures. Each time you call a @tf.function with a new argument signature, TensorFlow traces out a new graph just for that set of arguments. This new graph is then added as a "concrete function" to the callable. Thus, a saved model can be one or more subgraphs, each with a different signature.
A describing the functional . However, the material in this document lays out a snapshot of the existing state of things. Any links to code will include point-in-time revisions so as not to drift out of date. As with all non-documented implementation details, these details are subject to change in the future.
We’ll occasionally use the term “signatures” to talk about the general concept of describing function inputs (e.g. in the title of this document). In this sense, we will be referring not just to TensorFlow’s specific concept of signatures, but all of the ways TensorFlow defines and validates inputs to functions. Context should make the meaning clear.
What This Is Not About
This document is not intended to describe how signatures or functions work from a user perspective. It is intended for TensorFlow developers working on the internals of TensorFlow. Likewise, this document does not make a statement of the way things “should” be. It aims to simply document the way things are.
Overview of Signature Definitions
There are five protos that store definitions of function inputs in one manner or another. Their names and code locations, as well as their paths within the saved model proto, are as follows:
Proto messages, and their location in : -> SignatureDef.
On the loading side,SignatureDefs are essentially ignored. They are primarily used in v1 or C++, where the developer loading the model can inspect the returned SignatureDef protos directly. This allows them to use their desired signature name to lookup the placeholder and output names needed for execution.
These input and output names can then be passed into feeds and fetches when calling is one of the many types of . SavedFunctions are restored into a field.
SavedFunction’s main purpose is polymorphism. SavedFunctions support polymorphism by specifying a number of under the hood and is where the logic lives to find the appropriate concrete function.
This is invisible in the SavedModel proto, but these extra concrete functions are registered at call time in the runtime’s function library just as the other function library functions are.
The second purpose of SavedFunction is to update the FunctionSpec of all associated ConcreteFunctions using the FunctionSpec stored on the , and
SavedFunctions, they only reference a single are commonly attached to and accessed via the signatures map (i.e. the signatures attribute on the loaded object). The underlying concrete functions they modify, in this case, are (i.e. a dictionary of tensors). Similar to restored_function_body concrete functions, and other than restructuring the output, these concrete functions do nothing but call their associated concrete functions. SavedConcreteFunction
directly on the SavedObjectGraph. These objects reference a specific, already-registered concrete function -- the key in the map is that concrete function’s registered name.
These objects serve two purposes. The first is handling function "captures" via
the .
The second purpose, and similar to SavedFunction and SavedBareConcreteFunction, is modifying the existing concrete function’s FuncGraph structured inputs and outputs. This also is used for .
Example Walkthrough
A simple example may help illustrate all of this with more clarity. Let’s make a basic model and take a look at the subsequent generated proto to get a better feel for what’s going on.
Basic Model
class ExampleModel(tf.Module):
@tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
def capture_fn(self, x):
if not hasattr(self, 'weight'):
self.weight = tf.Variable(5.0, name='weight')
self.weight.assign_add(x * self.weight)
return self.weight
@tf.function
def polymorphic_fn(self, x):
return tf.constant(3.0) * x
model = ExampleModel()
model.polymorphic_fn(tf.constant(4.0))
model.polymorphic_fn(tf.constant([1.0, 2.0, 3.0]))
tf.saved_model.save(
model, "/tmp/example-model", signatures={'capture_fn': model.capture_fn})This model contains the basis for most of the complexity we’ll need to fully explore the intricacies of saving and signatures. This will allow us to look at functions with and without signatures, with and without captures, and with and without polymorphism.
Function with Captures
Let’s start by looking at our function with captures, capture_fn. We can see we have a concrete function defined in the function library, as expected:
A SavedConcreteFunction located in the concrete_functions map of the ObjectGraphDef |
Indeed, we can see bound_inputs refers to node 1, which is a SavedVariable with the name and dtype we expect:
A SavedBareConcreteFunction located in ObjectGraphDef.nodes |
As discussed above, we use the function spec and argument information to modify the underlying concrete function. But what’s up with the "__inference_signature_wrapper_68" name? And how does this fit in with the rest of the code?
First, note that this is the fifth (5) node in the node list. This will come up again shortly.
Let’s start by looking at the nodes list. If we start at the first node in the nodes list, we’ll see a "signatures" node attached as a child:
A SavedUserObject located in ObjectGraphDef.nodes |
Thus, when we access this function via model.signatures["capture_fn"], we will actually be calling into this intermediate signature wrapper function first.
And what does that function, "__inference_signature_wrapper_68", look like?
A SavedFunction located in ObjectGraphDef.nodes |
Again, the function spec is used to modify the function spec of our concrete function, "__inference_capture_fn_59". Notice also that concrete_functions here is a list. We only have one item right now, but this will come up again when we take a look at our polymorphic function example.
Now, we’ve fully mapped essentially everything needed for execution of this function, but we have one last thing to look at: SignatureDef. We’ve defined a signature, so we expect a SignatureDef to be defined:
A FunctionDef located in FunctionDefLibrary of MetaGraphDef.graph_def |
For one, we now have two concrete functions registered in the function library, each with slightly different input shapes.
We also have two SavedConcreteFunction modifiers:
A SavedFunction located in ObjectGraphDef.nodes |
The function spec here will be attached to both of these concrete functions at load time. When we call our SavedFunction, it will use the arguments we pass in to find the correct concrete function and execute it.
Next Steps
You should now be an expert on how functions and their signatures are saved at a code level. Remember, what's described in this blog post is how the code works right now. For updated code and examples in the future, see the official documentation on and itself, as well as a complete discussion of autograph.
And finally, if you do any exciting or useful protobuf surgery, share with us on Twitter. Thanks for reading this far!
SOCIAL SHARE CARD GENERATOR