Set static shapes in an existing tensorflow graph where dynamic shapes are used for input

Viewed 1911

I have been given a frozen (.pb) graph that has dynamic shapes for input (e.g. "None, None, None, 3" or "?x?x?x3"). I want to set these to a static shape (e.g. "1, 320, 320, 3"), however I am unsure of how I can change the shapes to the input placeholder that would apply the changes across all layers that follow. In this particular case, I do not have the code or ckpt files available so it is essential to make this work on a frozen (.pb) graph.

What have I already tried?

I have created a simple example code to make a simple graph and save it as frozen graph that can be used for testing different approaches. This graph is created as follows:

import tensorflow as tf

def simple_cnn_graph():
    graph = tf.Graph()
    with graph.as_default():
        input_layer = tf.placeholder(shape=[None, None, None, 3], dtype=tf.float32)
        conv1 = tf.layers.conv2d(
            inputs=input_layer,
            filters=32,
            kernel_size=[5, 5],
            padding="same",
            activation=tf.nn.relu)
        pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2, name='pool1')
        conv2 = tf.layers.conv2d(
            inputs=pool1,
            filters=16,
            kernel_size=[5, 5],
            padding="same",
            activation=tf.nn.relu)
        pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2, name='pool2')
    return graph, pool2

if __name__=='__main__':
    graph, output = simple_cnn_graph()
    with tf.Session(graph=graph) as sess:
        sess.run(tf.global_variables_initializer())
        graph_def = tf.graph_util.convert_variables_to_constants(sess, \
                            tf.get_default_graph().as_graph_def(), [output.name.split(':')[0]])

    frozen_file='./frozen.pb'
    with open(frozen_file, 'wb') as f:
        f.write(graph_def.SerializeToString())

    print([n.name for n in graph.as_graph_def().node])

I have tried two approaches:

1) I have tried using the transform_graph tool at: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/graph_transforms/README.md using strip_unused_nodes, however it does not work as I do not have any tensors that are converted to placeholders.

2) I have had some success following the comment in the link: https://github.com/tensorflow/tensorflow/issues/5680#issuecomment-405128390 Where I was able to use the tf.import_graph_def's input_map to map a new placeholder, however I am looking for a much simpler and generalizable solution that can be applied to any such frozen network graph in future (something similar to transform_graph for example). Below is the code I have for using the tf.import_graph_def method

import tensorflow as tf

def load_frozen_graph(frozen_file='frozen.pb'):
    graph = tf.Graph()
    with graph.as_default():
      od_graph_def = tf.GraphDef()
      with tf.gfile.GFile(frozen_file, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')
    return graph

graph = load_frozen_graph('./frozen.pb')

print('Tensor shapes before import map')
input_tensor = graph.get_tensor_by_name('Placeholder:0')
print(input_tensor)
output_tensor = graph.get_tensor_by_name('pool2/MaxPool:0')
print(output_tensor)

new_graph = tf.Graph()
with new_graph.as_default():
    new_input = tf.placeholder(dtype=tf.float32, shape=[1, 320, 320, 3], name='Placeholder')
    tf.import_graph_def(graph.as_graph_def(), name='', input_map={'Placeholder': new_input})

print('Tensor shapes after import map')
input_tensor = new_graph.get_tensor_by_name('Placeholder:0')
print(input_tensor)
output_tensor = new_graph.get_tensor_by_name('pool2/MaxPool:0')
print(output_tensor)

where the printed output is:

Tensor shapes before import map
Tensor("Placeholder:0", shape=(?, ?, ?, 3), dtype=float32)
Tensor("pool2/MaxPool:0", shape=(?, ?, ?, 16), dtype=float32)
Tensor shapes after import map
Tensor("Placeholder:0", shape=(1, 320, 320, 3), dtype=float32)
Tensor("pool2/MaxPool:0", shape=(1, 80, 80, 16), dtype=float32)

I will really appreciate if someone can point me in the right direction or correct me if I have made any mistake in the above code/post or understanding anything wrong about tf shapes.

0 Answers
Related