TensorFlow: Saving py_func to .pb file

Question

I try to build a tensorflow model - where i use the tf.py_func to create a part of the code in ordinary python code. The problem is that when I save the model to a .pb file, the .pb file itself is very small and does not include the py_func:0 tensor. When I try to load and run the model from the .pb file I get this error: get ValueError: callback pyfunc_0 is not found.

It works when I dont save and load as a .pb file

Is anyone able ton help. This is super important to me and have given me a couple of sleepless nights.

model_version = "465555564"
tensorboard = TensorBoard(log_dir='./logs', histogram_freq = 0, write_graph = True, write_images = False)

sess = tf.Session()
K.set_session(sess)
K.set_learning_phase(0)

def my_func(x):
    some_function

input = tf.placeholder(tf.float32)
y = tf.py_func(my_func, [input], tf.float32)

prediction_signature = tf.saved_model.signature_def_utils.predict_signature_def({"inputs": input}, {"prediction": y})
builder = saved_model_builder.SavedModelBuilder('./'+model_version)
legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
builder.add_meta_graph_and_variables(
      sess, [tag_constants.SERVING],
      signature_def_map={
           signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:prediction_signature,
      },
      legacy_init_op=legacy_init_op)

builder.save()

This is one of it's limitations. "The body of the function (i.e. func) will not be serialized in a GraphDef" — Sharky, May 21 '19 at 08:07
What I want to do is to load a .pickle file containing a "Random Forrest Model". That model I want to run inside Tensorflow — smallbirds, May 21 '19 at 08:32
You can try using autograph, https://www.tensorflow.org/guide/autograph — Sharky, May 21 '19 at 08:35
There is an implementation of random forests in TensorFlow, see [TensorForest](https://github.com/tensorflow/tensorflow/tree/r1.13/tensorflow/contrib/tensor_forest). Note this is currently under `tf.contrib`, which [will disappear in TensorFlow 2.x](https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md). See the RFC [TensorForest Estimator](https://github.com/tensorflow/community/blob/master/rfcs/20180626-tensor-forest.md) for the accepted plans for `tensor_forest`. — jdehesa, May 21 '19 at 09:00

score 4 · Accepted Answer · answered Jul 25 '19 at 18:16

There is a way to save TF models with tf.py_func, but you have to do it without using a SavedModel.

TF has 2 levels of model saving: checkpoints and SavedModels. See this answer for more details, but to quote it here:

A checkpoint contains the value of (some of the) variables in a TensorFlow model. It is created by a Saver. To use a checkpoint, you need to have a compatible TensorFlow Graph, whose Variables have the same names as the Variables in the checkpoint.

SavedModel is much more comprehensive: It contains a set of Graphs (MetaGraphs, in fact, saving collections and such), as well as a checkpoint which is supposed to be compatible with these Graphs, and any asset files that are needed to run the model (e.g. Vocabulary files). For each MetaGraph it contains, it also stores a set of signatures. Signatures define (named) input and output tensors.

The tf.py_func op cannot be saved with a SavedModel (noted on this page in the docs), which is what you tried to do here. There is a good reason for this. SavedModels are supposed to be totally independent from the original code, able to be loaded in any other language that can deserialize it. This allows the models to be loaded by things like ML Engine, which is probably written in C++ or something like that. The problem is that it cannot serialize arbitrary Python code, so py_func is a no-go.

You can work around this by using checkpoints, as long as you are okay with staying in Python. You will not get the independence that SavedModels provide. You can save a checkpoint after training with a tf.train.Saver, and then in a new Session, re-build the whole graph and load it with that Saver. There is even a way to use that code in ML Engine, which used to be exclusively for SavedModels. You can use custom prediction routines to side-step the need for a SavedModel.

More info on saving/restoring models in the docs.

TensorFlow: Saving py_func to .pb file

1 Answers1