How to wrap a frozen Tensoflow graph in a Keras Lambda layer in TF2?

Question

This question is related to this question, which provides a solution that works in Tensorflow 1.15, but doesn't work anymore in TF2

I'm taking part of the code from that question and adapting it slightly (removed the frozen model's multiple inputs and, with it, the need for nest).

Note: I'm separating the code in blocks, but they're meant to be run as on file (i.e., I won't repeat the unnecessary imports in each block)

First, we generate a frozen graph to use as dummy test network:

import numpy as np
import tensorflow.compat.v1 as tf

def dump_model():
    with tf.Graph().as_default() as gf:
        x = tf.placeholder(tf.float32, shape=(None, 123), name='x')
        c = tf.constant(100, dtype=tf.float32, name='C')
        y = tf.multiply(x, c, name='y')
        z = tf.add(y, x, name='z')
        with tf.gfile.GFile("tmp_net.pb", "wb") as f:
            raw = gf.as_graph_def().SerializeToString()
            print(type(raw), len(raw))
            f.write(raw)

dump_model()

Then, we load the frozen model and wrap it in a Keras Model:

persisted_sess = tf.Session()
with tf.Session().as_default() as session:
    with tf.gfile.FastGFile("./tmp_net.pb",'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        persisted_sess.graph.as_default()
        tf.import_graph_def(graph_def, name='')
        print(persisted_sess.graph.get_name_scope())
        for i, op in enumerate(persisted_sess.graph.get_operations()):
            tensor = persisted_sess.graph.get_tensor_by_name(op.name + ':0')
            print(i, '\t', op.name, op.type, tensor)
        x_tensor = persisted_sess.graph.get_tensor_by_name('x:0')
        y_tensor = persisted_sess.graph.get_tensor_by_name('y:0')
        z_tensor = persisted_sess.graph.get_tensor_by_name('z:0')

from tensorflow.compat.v1.keras.layers import Lambda, InputLayer
from tensorflow.compat.v1.keras import Model
from tensorflow.python.keras.utils import layer_utils

input_x = InputLayer(name='x', input_tensor=x_tensor)
input_x.is_placeholder = True
output_y = Lambda(lambda x: y_tensor, name='output_y')(input_x.output)
output_z = Lambda(lambda x_b: z_tensor, name='output_z')(input_x.output)

base_model_inputs = layer_utils.get_source_inputs(input_x.output)
base_model = Model(base_model_inputs, [output_y, output_z])

Finally, we run the model on some random data and verify that it runs without errors:

y_out, z_out = base_model.predict(np.ones((3, 123), dtype=np.float32))
y_out.shape, z_out.shape

In Tensorflow 1.15.3, the output of the above is ((3, 123), (3, 123)), however, if I run the same code in Tensorflow 2.1.0, the first two blocks run without a problem, but then the third fails with:

TypeError: An op outside of the function building code is being passed
a "Graph" tensor. It is possible to have Graph tensors
leak out of the function building context by including a
tf.init_scope in your function building code.
For example, the following function will fail:
  @tf.function
  def has_init_scope():
    my_constant = tf.constant(1.)
    with tf.init_scope():
      added = my_constant * 2
The graph tensor has name: y:0

The error seems related to Tensorflow's automatic "compilation" and optimization of functions, but I don't know how to interpret it, what the source of the error is, or how to resolve.

What is the correct way to wrap the frozen model in Tensorflow 2?

jdehesa · Accepted Answer · 2020-07-09T14:43:08.583

I can run your whole example fine in 2.2.0 like this.

import tensorflow as tf
from tensorflow.core.framework.graph_pb2 import GraphDef
import numpy as np

with tf.Graph().as_default() as gf:
    x = tf.compat.v1.placeholder(tf.float32, shape=(None, 123), name='x')
    c = tf.constant(100, dtype=tf.float32, name='c')
    y = tf.multiply(x, c, name='y')
    z = tf.add(y, x, name='z')
    with open('tmp_net.pb', 'wb') as f:
        f.write(gf.as_graph_def().SerializeToString())

with tf.Graph().as_default():
    gd = GraphDef()
    with open('tmp_net.pb', 'rb') as f:
        gd.ParseFromString(f.read())
    x, y, z = tf.graph_util.import_graph_def(
        gd, name='', return_elements=['x:0', 'y:0', 'z:0'])
    del gd
    input_x = tf.keras.layers.InputLayer(name='x', input_tensor=x)
    input_x.is_placeholder = True
    output_y = tf.keras.layers.Lambda(lambda x: y, name='output_y')(input_x.output)
    output_z = tf.keras.layers.Lambda(lambda x: z, name='output_z')(input_x.output)

    base_model_inputs = tf.keras.utils.get_source_inputs(input_x.output)
    base_model = tf.keras.Model(base_model_inputs, [output_y, output_z])

    y_out, z_out = base_model.predict(np.ones((3, 123), dtype=np.float32))
    print(y_out.shape, z_out.shape)
    # (3, 123) (3, 123)

The "trick" is to wrap the model construction within a with tf.Graph().as_default(): block, which will ensure everything is created in graph mode within the same graph object.

However, it may be simpler to wrap the graph loading and computation within a @tf.function, which would avoid this kind of error and make the model construction more transparent:

import tensorflow as tf
from tensorflow.core.framework.graph_pb2 import GraphDef
import numpy as np

@tf.function
def my_model(x):
    gd = GraphDef()
    with open('tmp_net.pb', 'rb') as f:
        gd.ParseFromString(f.read())
    y, z = tf.graph_util.import_graph_def(
        gd, name='', input_map={'x:0': x}, return_elements=['y:0', 'z:0'])
    return [y, z]

x = tf.keras.Input(shape=123)
y, z = tf.keras.layers.Lambda(my_model)(x)
model = tf.keras.Model(x, [y, z])
y_out, z_out = model.predict(np.ones((3, 123), dtype=np.float32))
print(y_out.shape, z_out.shape)
# (3, 123) (3, 123)

Thanks! It seems the key here is having the graph be the default graph both at model creation and inference. In my actual code, I'm loading and building the model in a function. If I return both the graph and the model, and use `with graph.as_default():` before calling predict()`, it works! If I don't, though, I still have the same error as above. Thank you very much for this answer, it gave me just the right information I was missing (like the fact that I don't need `Session`s to deal with graphs...). — GPhilo, Jul 09 '20 at 14:39
As a possible improvement point, would you perhaps clarify what you changed from my code, so that future viewers may also benefit from this solution? — GPhilo, Jul 09 '20 at 14:40
@GPhilo Yes, thanks for the suggestion. I have also added another possible way to do the same with `@tf.function`, which I think may be cleaner. The only minor downside is that you have to specify in advance the size of the input although you should probably know the size of the inputs and outputs for the model you're loading anyway. — jdehesa, Jul 09 '20 at 14:44
Your second option looks very much like what I had hoped would be possible! When you say you need to know the size of the input, does that have to be fully known, or can it have `None`s as well? (I'm working with a detection network, which unfortunately can - and needs to - accept variable-shaped input, so the input shape is `[None, None, None, 3]`) — GPhilo, Jul 09 '20 at 14:47
Ok, I answered that myself: It's fine to pass `None`s in the `Input` definition. I have, however, other problems with the NMS operation and something about inputs from different frames. I remember having this problem a while ago, but I can't remember the solution. It doesn't matter, though: your first approach works perfectly! — GPhilo, Jul 09 '20 at 15:01
@GPhilo Ah, yes, you can have `None` (in fact, if you pass `shape=(None,)` in the example it also works). I meant to say, you need to give the model `Input` a shape that is compatible with the shape of the input in the model. — jdehesa, Jul 09 '20 at 15:01

score 2 · Answer 2 · answered Jul 09 '20 at 14:47

Another possible way to do this would be

import tensorflow as tf

input_layer = tf.keras.Input(shape=[123])
keras_graph = input_layer.graph

with keras_graph.as_default():
    with tf.io.gfile.GFile('tmp_net.pb', 'rb') as f:
        graph_def = tf.compat.v1.GraphDef()
        graph_def.ParseFromString(f.read())

    tf.graph_util.import_graph_def(graph_def, name='', input_map={'x:0': input_layer})
    
    
y_tensor = keras_graph.get_tensor_by_name('y:0')
z_tensor = keras_graph.get_tensor_by_name('z:0')

base_model = tf.keras.Model(input_layer, [y_tensor, z_tensor])

And then

y_out, z_out = base_model.predict(tf.ones((3, 123), dtype=tf.float32))
print(y_out.shape, z_out.shape)
# (3, 123) (3, 123)

I didn't know that `Input` layers had a `graph` attribute, thanks! This looks very interesting too! — GPhilo, Jul 09 '20 at 15:04
This wouldn't work for newer TensorFlow versions, at least it's not working for TF v2.9.1. as: `AttributeError: 'KerasTensor' object has no attribute 'graph'` — Kuzman Belev, Jul 12 '23 at 11:03

How to wrap a frozen Tensoflow graph in a Keras Lambda layer in TF2?

2 Answers2