TensorFlow Graph to Keras Model?

Question

Is it possible to define a graph in native TensorFlow and then convert this graph to a Keras model?

---

My intention is simply combining (for me) the best of the two worlds.

I really like the Keras model API for prototyping and new experiments, i.e. using the awesome multi_gpu_model(model, gpus=4) for training with multiple GPUs, saving/loading weights or whole models with oneliners, all the convenience functions like .fit(), .predict(), and others.

However, I prefer to define my model in native TensorFlow. Context managers in TF are awesome and, in my opinion, it is much easier to implement stuff like GANs with them:

with tf.variable_scope("Generator"):
    # define some layers
with tf.variable_scope("Discriminator"):
    # define some layers

# model losses
G_train_op = ...AdamOptimizer(...)
    .minimize(gloss,
    var_list=tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, 
                               scope="Generator")
D_train_op = ...AdamOptimizer(...)
    .minimize(dloss, 
    var_list=tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, 
                               scope="Discriminator")

Another bonus is structuring the graph this way. In TensorBoard debugging complicated native Keras models are hell since they are not structured at all. With heavy use of variable scopes in native TF you can "disentangle" the graph and look at a very structured version of a complicated model for debugging.

By utilizing this I can directly setup custom loss function and do not have to freeze anything in every training iteration since TF will only update the weights in the correct scope, which is (at least in my opinion) far easier than the Keras solution to loop over all the existing layers and set .trainable = False.

TL;DR:

Long story short: I like the direct access to everything in TF, but most of the time a simple Keras model is sufficient for training, inference, ... later on. The model API is much easier and more convenient in Keras.

Hence, I would prefer to set up a graph in native TF and convert it to Keras for training, evaluation, and so on. Is there any way to do this?

Answer 1

I don't think it is possible to create a generic automated converter for any TF graph, that will come up with a meaningful set of layers, with proper namings etc. Just because graphs are more flexible than a sequence of Keras layers.

However, you can wrap your model with the Lambda layer. Build your model inside a function, wrap it with Lambda and you have it in Keras:

def model_fn(x):
    layer_1 = tf.layers.dense(x, 100)
    layer_2 = tf.layers.dense(layer_1, 100)
    out_layer = tf.layers.dense(layer_2, num_classes)
    return out_layer

model.add(Lambda(model_fn))

That is what sometimes happens when you use multi_gpu_model: You come up with three layers: Input, model, and Output.

Keras Apologetics

However, integration between TensorFlow and Keras can be much more tighter and meaningful. See this tutorial for use cases.

For instance, variable scopes can be used pretty much like in TensorFlow:

x = tf.placeholder(tf.float32, shape=(None, 20, 64))
with tf.name_scope('block1'):
    y = LSTM(32, name='mylstm')(x)

The same for manual device placement:

with tf.device('/gpu:0'):
    x = tf.placeholder(tf.float32, shape=(None, 20, 64))
    y = LSTM(32)(x)  # all ops / variables in the LSTM layer will live on GPU:0

Custom losses are discussed here: Keras: clean implementation for multiple outputs and custom loss functions?

This is how my model defined in Keras looks in Tensorboard:

So, Keras is indeed only a simplified frontend to TensorFlow so you can mix them quite flexibly. I would recommend you to inspect source code of Keras model zoo for clever solutions and patterns that allows you to build complex models using clean API of Keras.

Answer 2

You can insert TensorFlow code directly into your Keras model or training pipeline! Since mid-2017, Keras has fully adopted and integrated into TensorFlow. This article goes into more detail.

This means that your TensorFlow model is already a Keras model and vice versa. You can develop in Keras and switch to TensorFlow whenever you need to. TensorFlow code will work with Keras APIs, including Keras APIs for training, inference and saving your model.