How to import a model saved as a TFHub module, into lucid?

[mencia]

Dear all,

I want to visualize a VAE model that is saved as a TFHub module, in directory path. I can load it using hub.Module:

import tensorflow_hub as hub module = hub.Module(path)

Relevant information about the model:

module.variables [<tf.Variable 'module_1/decoder/conv2d_transpose/bias:0' shape=(64,) dtype=float32>, <tf.Variable 'module_1/decoder/conv2d_transpose/kernel:0' shape=(4, 4, 64, 64) dtype=float32>, <tf.Variable 'module_1/decoder/conv2d_transpose_1/bias:0' shape=(32,) dtype=float32>, <tf.Variable 'module_1/decoder/conv2d_transpose_1/kernel:0' shape=(4, 4, 32, 64) dtype=float32>, <tf.Variable 'module_1/decoder/conv2d_transpose_2/bias:0' shape=(32,) dtype=float32>, <tf.Variable 'module_1/decoder/conv2d_transpose_2/kernel:0' shape=(4, 4, 32, 32) dtype=float32>, <tf.Variable 'module_1/decoder/conv2d_transpose_3/bias:0' shape=(1,) dtype=float32>, <tf.Variable 'module_1/decoder/conv2d_transpose_3/kernel:0' shape=(4, 4, 1, 32) dtype=float32>, <tf.Variable 'module_1/decoder/dense/bias:0' shape=(256,) dtype=float32>, <tf.Variable 'module_1/decoder/dense/kernel:0' shape=(10, 256) dtype=float32>, <tf.Variable 'module_1/decoder/dense_1/bias:0' shape=(1024,) dtype=float32>, <tf.Variable 'module_1/decoder/dense_1/kernel:0' shape=(256, 1024) dtype=float32>, <tf.Variable 'module_1/encoder/e1/bias:0' shape=(32,) dtype=float32>, <tf.Variable 'module_1/encoder/e1/kernel:0' shape=(4, 4, 1, 32) dtype=float32>, <tf.Variable 'module_1/encoder/e2/bias:0' shape=(32,) dtype=float32>, <tf.Variable 'module_1/encoder/e2/kernel:0' shape=(4, 4, 32, 32) dtype=float32>, <tf.Variable 'module_1/encoder/e3/bias:0' shape=(64,) dtype=float32>, <tf.Variable 'module_1/encoder/e3/kernel:0' shape=(2, 2, 32, 64) dtype=float32>, <tf.Variable 'module_1/encoder/e4/bias:0' shape=(64,) dtype=float32>, <tf.Variable 'module_1/encoder/e4/kernel:0' shape=(2, 2, 64, 64) dtype=float32>, <tf.Variable 'module_1/encoder/e5/bias:0' shape=(256,) dtype=float32>, <tf.Variable 'module_1/encoder/e5/kernel:0' shape=(1024, 256) dtype=float32>, <tf.Variable 'module_1/encoder/log_var/bias:0' shape=(10,) dtype=float32>, <tf.Variable 'module_1/encoder/log_var/kernel:0' shape=(256, 10) dtype=float32>, <tf.Variable 'module_1/encoder/means/bias:0' shape=(10,) dtype=float32>, <tf.Variable 'module_1/encoder/means/kernel:0' shape=(256, 10) dtype=float32>]

module.get_input_info_dict(signature='gaussian_encoder') module.get_input_info_dict(signature='decoder') module.get_input_info_dict(signature='reconstructions') {'images': <hub.ParsedTensorInfo shape=(?, 64, 64, 1) dtype=float32 is_sparse=False>} {'latent_vectors': <hub.ParsedTensorInfo shape=(?, 10) dtype=float32 is_sparse=False>} {'images': <hub.ParsedTensorInfo shape=(?, 64, 64, 1) dtype=float32 is_sparse=False>}

module.get_output_info_dict(signature='gaussian_encoder') module.get_output_info_dict(signature='decoder') module.get_output_info_dict(signature='reconstructions') {'mean': <hub.ParsedTensorInfo shape=(?, 10) dtype=float32 is_sparse=False>, 'logvar': <hub.ParsedTensorInfo shape=(?, 10) dtype=float32 is_sparse=False>} {'images': <hub.ParsedTensorInfo shape=(?, 64, 64, 1) dtype=float32 is_sparse=False>} {'images': <hub.ParsedTensorInfo shape=(?, 64, 64, 1) dtype=float32 is_sparse=False>}

Importing model into Lucid:

The way to do it is:

from lucid.modelzoo.vision_models import Model

with tf.Graph().as_default() as graph, tf.Session() as sess:
      images = tf.placeholder("float32", [None, 224, 224, 3], name="input")

     # <Code to construct & load your model inference graph goes here>
     # ...

    Model.save(
      "saved_model.pb",
      image_shape=[W, W, 3],
      input_name='input',
      output_names=['Softmax'],
      image_value_range=[0,1],
    )

Now, I don't know what I should do in <Code to construct & load your model inference graph goes here>. Let us say I want to visualize a certain channel in one of the encoder layers. How should I construct & load my model inference graph?

[colah]

I've never used tfhub, but you probably just do something like module(images)?

Unfortunately, this is more of a tfhub qusestion than a lucid question.

[mencia]

Thank you very much for your answer @colah ! Let me reformulate the question in a way where the tfhub part is sorted out, and we can focus on the lucid part. I already have an inference graph of my tfhub VAE model. I select the encoder part of the model, such that the input is an image and the output the logvar and mean. Here is the code that does it:

import tensorflow as tf
import tensorflow_hub as hub 
path = '../0/model/tfhub'

module = hub.Module(path)
outputs = module(dict(images=dataset), signature='gaussian_encoder',as_dict=True)
logvar = outputs["logvar"]
mean = outputs["mean"]

sess = tf.Session()
sess.run(tf.global_variables_initializer())
sess.run([logvar, mean])

This outputs the encoded images, i.e., logvar and mean vectors. Here dataset is a numpy array with shape (2, 64, 64, 1). Now, the way I try to import this into lucid is:

from lucid.modelzoo.vision_models import Model

with tf.Graph().as_default() as graph, tf.Session() as sess:

    images = tf.placeholder("float32", [None, 64, 64, 1], name="images")
    module = hub.Module(path)
    outputs = module(dict(images=images), signature='gaussian_encoder',as_dict=True)
    logvar = outputs['logvar']
    mean = outputs['mean']

    Model.save("saved_model.pb",  
     input_name='images', 
     output_names=['logvar','mean'], 
     image_shape=[64, 64, 1],
     image_value_range=[0,1])

I get the following error:

INFO:tensorflow:Saver not created because there are no variables in the graph to restore

INFO:tensorflow:Saver not created because there are no variables in the graph to restore

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-7-87ae8bf407d3> in <module>
     13      output_names=['logvar','mean'],
     14      image_shape=[64, 64, 1],
---> 15      image_value_range=[0,1])

~/<...>/env_py36/lib/python3.6/site-packages/lucid/modelzoo/vision_base.py in save(save_url, input_name, output_names, image_shape, image_value_range)
    309     }
    310 
--> 311     graph_def = model_util.frozen_default_graph_def([input_name], output_names)
    312     model_util.infuse_metadata(graph_def, metadata)
    313     save(graph_def, save_url)

~/<...>/env_py36/lib/python3.6/site-packages/lucid/modelzoo/util.py in frozen_default_graph_def(input_node_names, output_node_names)
     73       input_graph_def, protected_nodes=(output_node_names + input_node_names)
     74   )
---> 75   pruned_graph = tf.graph_util.extract_sub_graph(pruned_graph, output_node_names)
     76 
     77   # remove explicit device assignments

~/<...>/env_py36/lib/python3.6/site-packages/tensorflow/python/util/deprecation.py in new_func(*args, **kwargs)
    322               'in a future version' if date is None else ('after %s' % date),
    323               instructions)
--> 324       return func(*args, **kwargs)
    325     return tf_decorator.make_decorator(
    326         func, new_func, 'deprecated',

~/<...>/env_py36/lib/python3.6/site-packages/tensorflow/python/framework/graph_util_impl.py in extract_sub_graph(graph_def, dest_nodes)
    180   name_to_input_name, name_to_node, name_to_seq_num = _extract_graph_summary(
    181       graph_def)
--> 182   _assert_nodes_are_present(name_to_node, dest_nodes)
    183 
    184   nodes_to_keep = _bfs_for_reachable_nodes(dest_nodes, name_to_input_name)

~/<...>/env_py36/lib/python3.6/site-packages/tensorflow/python/framework/graph_util_impl.py in _assert_nodes_are_present(name_to_node, nodes)
    135   """Assert that nodes are present in the graph."""
    136   for d in nodes:
--> 137     assert d in name_to_node, "%s is not in graph" % d
    138 
    139 

AssertionError: logvar is not in graph

Do you have any suggestions for how to overcome this?

[mencia]

Solution:

output_names should be the name of one of the nodes of the graph. One can access the nodes in the following way:

with tf.Graph().as_default() as graph, tf.Session() as sess:

    images = tf.placeholder("float32", [None, 64, 64, 1], name="images")
    module = hub.Module(path)
    outputs = module(dict(images=images), signature='gaussian_encoder',as_dict=True)
    print('number of nodes =',len(tf.get_default_graph().as_graph_def().node))
    print('node names:')
    [print(n.name) for n in tf.get_default_graph().as_graph_def().node]

The following code saves the tfhub model:

from lucid.modelzoo.vision_models import Model
with tf.Graph().as_default() as graph, tf.Session() as sess:

    images = tf.placeholder("float32", [None, 64, 64, 1], name="images")
    module = hub.Module(path)
    outputs = module(dict(images=images), signature='gaussian_encoder',as_dict=True)

    sess.run(tf.global_variables_initializer())

    Model.save("saved_model.pb",  
     input_name='images', 
     output_names=[graph.as_graph_def().node[-1].name], 
     image_shape=[64, 64, 1],
     image_value_range=[0,1])

Note that I chose the last node as my output. The output nodes will be kept unconditionally when the nodes that are not needed for inference are pruned out. Furthermore, the saved graph will be able to reach these nodes.

It remains for me to figure out how to extract layer_names from saved_model.pb to visualize layers:

model = Model.load("saved_model.pb")
render.render_vis(model, "layer_name:0")