Importing models of custom tf.Module classes

[Leandro Graciá Gil]

Hi,

I'm doing some experiments with tf.Module as part of migrating our code to TF 2.0. In the saved model guide there's information about how to save modules and how to provide signatures that allow them to be loaded back.

As expected, the python program itself is not saved in the model and as such it's not automatically restored. This makes storing custom tf.Module objects partially one-directional: you can save them but you don't get back the same when you load them (you recover your exported TF functions and data, but lose your python functions and data). This fine from automatic model loading since we're talking about loading custom classes TF doesn't know about, but it should still be possible for my own custom tf.Module class to be able to fully restore itself from a saved model.

My question is, what would be the correct / recommended way to restore an object of my own custom tf.Module class using the results of tf.saved_model.load? This of course assuming I can somehow recover any other required non-TF data for construction and such.

Thanks,

Leandro

PS: I'm aware that Keras models can be restored directly to python objects, but we're not planning to use the Keras API at the moment.

[Martin Wicke]

+Kathy Wu 

This isn't something tf.module supports. It is intentionally minimal. The reason Layer supports restoring is because of the from_config/get_config methods, complexity of everybody needs. 

What is the reason you cannot use Layer (without necessarily using Model or other Keras utilities?)

Martin

--
You received this message because you are subscribed to the Google Groups "TensorFlow Developers" group.
To unsubscribe from this group and stop receiving emails from it, send an email to developers+...@tensorflow.org.
To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/developers/CALWiF1syWS%3DEw3E-irYnBEu%2BzJG%3D5qz9CtkVosWoXK%2Bn8EFoAA%40mail.gmail.com.

[Leandro Gracia Gil]

On 18/09/2019 23:21, Martin Wicke wrote:
> What is the reason you cannot use Layer (without necessarily using
> Model or other Keras utilities?)
>

We’re trying to migrate the main features of our framework built on top
of TF 1.x to TF 2.0. In particular, we have a generic model abstraction
that allows recursive composability, including the ability to reuse
inner submodels at will and to compute sets of trainable variables. In
fact, it’s remarkably similar to what tf.Module looks like.

However, our approach is fully deserializable from your custom python
class, which is the bit I’m missing here. Usually our model classes
would have both python methods and TF-based methods that would
correspond to the ones decorated with tf.function. We would achieve
automatic serialization by registering our custom type and providing a
couple of methods to read and write information into/from the object's
name scope in the TF graph, including any required python information
that would be saved as constants. With this we could easily restore the
python object of one of our custom models from a graphdef file and its
name scope, for example.

I should be able to get something similar working on top of TF 2.0, but
without graphdefs it looks like tf.saved_model is the way to go for
serialization. Without a clear way to extend saving and loading in
tf.Modules, I guess I could try to somehow wrap around what
tf.saved_model.load returns, but that sounds like the kind of thing that
can cause trouble with tf.function decorations, the existing tf.Module
code and such.

So, any advice on what could be a reasonable approach to do this or what
kind of things I should avoid? Maybe there's something similar to
from_config/get_config I can do?

Leandro

[Martin Wicke]

The question is, why do you not want to use the existing config methods? I.e., change nothing to what you're doing now, but instead of tf.Module, use tf.Layer. 

I am not sure this would solve your problem, but I'd like to understand why it doesn't.

[Leandro Gracia Gil]

Thanks for pointing to the existing config methods. I wasn't fully aware of how they work. Looking at them, it seems they might work as a base for what I need.

There are some differences, though. In our approach we'd be saving the config together with the model data, which I guess I can do by simply saving a new file with the config in the model folder.

Another difference would be that our graphdefs could store any number of custom objects, and these would be findable in the graph by their name scope. Here it's more like if each object corresponds to a separate saved model. This can get tricky if we have submodules that are shared between other submodules (a DAG vs a tree), because instead of simply referring to them we'll probably have them saved multiple times and loaded as separate entities.

I'll experiment with the existing config methods. Meanwhile, any advice on how to solve the above issue would be most welcome.

Thanks!
Leandro

[Martin Wicke]

Layers with sublayers are saved and restored properly, so that should work in your case as well.

I will warn that relying on name scopes and generally, names are not something I would recommend, it's very brittle. 

[Leandro Gracia Gil]

Yes, I've noticed that recursive modules (and I guess layers) are saved and restored automatically, but that raises a question. Let's say we have this layer structure.

Layer A: contains sublayers B and C.
Layer B: contains sublayer D.
Layer C: contains sublayer D.
Layer D: no sublayers.
Layer E: contains sublayer D.

If I try to use tf.save_model.save with layers A and E, won't that save at least 2 different copies of layer D? Is there any way I can avoid this when saving and loading?

[Martin Wicke]

You might restore two copies of D if you save A and E separately, but if you had a Layer F which contains both, you should save only one. 

Kathy, can you confirm what would happen here?

[Kathy Wu]

Confirming Martin's words - if you save the layers separately, then you'll restore two copies of layer D. Saving both layers at once under a single root layer or module would restore the same copy D.

[Leandro Graciá Gil]

Sounds good. Thanks!

[Leandro Graciá Gil]

After looking further into this I have couple of questions about Layer.from_config:

1. How can Layer.from_config be used to create a python object that uses the TF variables/functions/etc. loaded from a SavedModel? The docs claim it provides serialization, but it's not that clear how the TF graph elements are reused under the hood. For all I know it could be creating a new identical layer with new weights.
   
   
2. If the SavedModel contained sublayers, how can the above be done for each sublayer recursively?

Basically, I'm trying to somehow save the layer configs together with the SavedModel data in order to easily restore it as a python object. What I'm finding is that first I need a way to associate each layer to its config (since there might be an arbitrary number of sublayers in the SavedModel, each with its own config). Then I need to be able to run Layer.from_config on each of the sublayers to get the hierarchy of python objects back, keeping in mind that a same sublayer might be referenced multiple times by others (DAG case).

Any suggestions about how this could be done?

Thanks!

[Tom Hennigan]

I'm not familiar with Keras so will let someone else answer re that.

Personally in this case (since you want the Python object not some proxy with an equivalent TF graph) I would suggest keeping your Python code around (e.g. make use of a git repo to version your experiment code and deps) and using a tf.train.Checkpoint to save/restore TensorFlow state. We have an example of checkpointing here and I think all you would need to do is replace "create_my_sonnet_module()" with your code to construct the initial state of your network. TensorFlow's checkpointing supports restore on create, so as your tf.Variables are created they will inherit the correct initial value. AFAIK this will work with a completely arbitrary DAG of modules (I would suggest testing though).

You could also consider using pickle if you want a zero config way to serialize arbitrary Python objects (tf.Variable and tf.Module support pickle), but for this to work you will still need to version your code (pickle is very flaky even after minor changes) so I think the first option is best.

--
You received this message because you are subscribed to the Google Groups "TensorFlow Developers" group.
To unsubscribe from this group and stop receiving emails from it, send an email to developers+...@tensorflow.org.

To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/developers/CALWiF1vJPgxr%2BLS6W-0r3BCqC57_F6CH8o8XYvUAbOo9i%2BUnAw%40mail.gmail.com.

[Kathy Wu]

Are you asking about how `tf.keras.models.load_model` works, or how you could implement a loading function that uses Layer.from_config?

Here's a brief summary of how SavedModel saving/loading works in Keras: The model and all of its layers are saved individually with their own list of variables, call functions, etc. When loaded, each object is restored into custom layer/model objects that use the variables and functions from the SavedModel. In other words, what you're getting is not an instance of the same class you're saving but a python object that acts similarly. 

The layer hierarchy in the original model will be restored in the loaded object. It seems like `tf.keras.models.load_model` fits your needs, but let us know if it doesn't.

On Sunday, September 22, 2019 at 10:46:22 AM UTC-7, Leandro Graciá Gil wrote:

After looking further into this I have couple of questions about Layer.from_config:

1. How can Layer.from_config be used to create a python object that uses the TF variables/functions/etc. loaded from a SavedModel? The docs claim it provides serialization, but it's not that clear how the TF graph elements are reused under the hood. For all I know it could be creating a new identical layer with new weights.
   
   
2. If the SavedModel contained sublayers, how can the above be done for each sublayer recursively?

Basically, I'm trying to somehow save the layer configs together with the SavedModel data in order to easily restore it as a python object. What I'm finding is that first I need a way to associate each layer to its config (since there might be an arbitrary number of sublayers in the SavedModel, each with its own config). Then I need to be able to run Layer.from_config on each of the sublayers to get the hierarchy of python objects back, keeping in mind that a same sublayer might be referenced multiple times by others (DAG case).

Any suggestions about how this could be done?

Thanks!

[Leandro Graciá Gil]

Thank you everyone for your suggestions. I’ve been giving them a try, digging into the source code, and exploring a few other ideas. Let me share what I found to be the most reasonable approach for my use case.

It seems that the simplest for me is to be based on Keras models, even if I’m not particularly in need of the Keras high level APIs. In particular, I’m creating a couple of base classes for custom layers and models. Here’s what I do:

For custom layers:

- I use python 3.6+ __init_subclass_ to automatically register any derived classes using tf.keras.utils.get_custom_objects(). With that, get_config and from_config can restore custom layer types. Note that this does not restore their weights, instead this is done at model level.

For custom models:

- To serialize I first use tf.keras.Model.save with serving_only=True (which is required for now). This creates a SavedModel that I can already use for serving if desired.

- Inside the saved model I create an “assets.extra” folder (as suggested here) and serialize two things: the model config (from a get_config that models must provide) and the input shape passed to the Model.build method.

- To deserialize, I first create a new python object using the saved config. Note that from_config needs to be overriden to just pass the arguments to the constructor, as tf.keras.Model.from_config seems to expect a custom config format for the functional API.

- I call build on the new python object with the saved input shape, so that any weights are created.

- I restore the saved model with tf.keras.models.load_model, then I basically do py_object.set_weights(saved_model.get_weights()).

This seems to work reasonably well and allows having something like model.save(path) and CustomModel.load(path) by just extending a base class, even if it has custom layers. However, there are a few things that don’t work.

- Because I didn’t call compile on the new python object, high level Keras APIs such as fit do not work. Custom train steps using gradient tapes do seem to work as usual, which is enough for my use case. Still, this could be fixed by saving the arguments to compile, but these will usually include optimizers and arbitrary python functions for losses and metrics, so it gets tricky to serialize.

- Using a functional keras model from inside a subclassed one does not work, although that’s likely not related to serialization but to eager inputs vs placeholder issues. This can be a problem if you need to use a third-party functional model, though, so any advice on this would be most appreciated.

Does this approach sound reasonable? Anything I should be aware of?

Thank you very much for your help.

Leandro

[Kathy Wu]

Hi Leandro,

It sounds like you might be using `tf.keras.experimental.export_saved_model`. This will be removed from the API in 2.1. Have you tried using `model.save(path, save_format='tf')`? It should be able to serialize custom layers/models without having to go through the extra registration steps. In addition, it will save your compiled parameters and will recompile the model when it is loaded.

-Kathy

[Leandro Graciá Gil]

Hi Kathy,

Sorry, I made a mistake in my description of serializing custom models. I originally tried using tf.keras.experimental.export_saved_model, but eventually discarded it in favor of model.save as you suggest. Therefore, the part mentioning the use of serving_only=True does not apply and should be ignored.

However, note that I found using model.save alone not to be enough for my purposes for 2 reasons:

1. For the subclassed model case, using tf.keras.models.load_model might recompile the model, but what I get it's not a python object of its original type. Since I'm building a new python object and setting its weights manually from the ones in the SavedModel I loaded, I'm also forced to call compile again on the new object even if only to set up its state. If I don't then I get errors like that the python object is missing its 'loss' attribute if I call its fit method.
   
   Similar to setting the weights, is there any way I can use what I get from tf.keras.models.load_model to avoid calling compile again on the new python object? I couldn't find any, but if there is one I'd love to know.
   
   
2. For the functional model case, registering custom layer types with tf.keras.utils.get_custom_objects() seems to be required. Otherwise you get a ValueError exception when deserializing the model that says "Unknown layer: (your custom layer type)".

Leandro