How to convert trained Keras model to a single TensorFlow .pb file and make prediction
You are going to learn step by step how to freeze and convert your trained Keras model into a single TensorFlow .pb file.
When compared to TensorFlow, Keras API might look less daunting and easier to work with, especially when you are doing quick experiments and build a model with standard layers. While TensorFlow is more versatile when you plan to deploy your model to different platforms across different programming languages. While there are many ways to convert a Keras model to its TenserFlow counterpart, I am going to show you one of the easiest when all you want is to make predictions with the converted model in deployment situations.
Here is the overview what will be covered.
- Keras to single TensorFlow .pb file
- Load .pb file with TensorFlow and make predictions.
- (Optional) Visualize the graph in a Jupyter notebook.
Source code for this post available on my GitHub.
Keras to TensorFlow .pb file
When you have trained a Keras model, it is a good practice to save it as a single HDF5 file first so you can load it back later after training.
Then you can load up the model and find the model’s input and output tensors’ names.
As you can see, our simple model has only single input and output, your model might have multiple inputs/outputs.
We keep track of their names since we are going to locate them by name in the converted TensorFlow graph during inference.
The first step is to get the computation graph of TensorFlow backend which represents the Keras model, where the forward pass and training related operations are included.
Then the graph will be converted to a GraphDef protocol buffer, after that it will be pruned so subgraphs that are not necessary to compute the requested outputs such as the training operations are removed. This step if refer to as freezing the graph.
The frozen_graph is a serialized GraphDef proto which we can use the following function call to save it as a single binary pb file.
Load .pb file and make predictions
Now we have everything we need to predict with the graph saved as one single .pb file.
To load it back, start a new session either by restarting the Jupyter Notebook Kernel or running in a new Python script.
The following several lines deserialize the GraphDef from .pb file and restore it as the default graph to current running TensorFlow session.
Locate the input tensor so we can feed it with some input data and grab the predictions from the output tensor, we are going to get them by name. The only difference is that all tensors’ names are prefixed with the string “import/” so the input tensor is now named "import/conv2d_1_input:0" and output tensor is "import/dense_2/Softmax:0".
To make a prediction, it can be as simple as,
The predictions is the softmax value in this case with shape (20, 10) representing 20 samples each has logits values for 10 classes.
If your model has multiple inputs/outputs, you can do something like this.
Visualize the graph in Notebook (optional)
Do you wonder what the model freezing step have done to your model, like what operations have been removed?
Let’s compare those two graphs side by side in the Jupyter Notebook by loading a minimal version of TensorBoard to show the graph structures.
Here I included a show_graph.py module allowing you to do so.
You can run this block twice, one after Keras model training/loading, one after loading&restoring the .pb file, and here is the results.
You can easily see the related training operations are removed in the frozen graph.
Conclusion and further reading
You have learned how to convert your Keras model into a TensorFlow .pb file for inference purpose only. Be sure to check out the source code for this post on my GitHub.
Here are some related resources you might find helpful.
A Tool Developer’s Guide to TensorFlow Model Files
Originally published at www.dlology.com.
