Skip to content

Instantly share code, notes, and snippets.

@dkurt
Last active August 15, 2018 07:35
Show Gist options
  • Save dkurt/f9bf5434f34b2c9d4913e8bf57854ba7 to your computer and use it in GitHub Desktop.
Save dkurt/f9bf5434f34b2c9d4913e8bf57854ba7 to your computer and use it in GitHub Desktop.

Let's discuss how to deploy deep learning models trained in TensorFlow framework both in TensorFlow and using OpenCV.

Requirements

  • TensorFlow
  • TensorBoard (optional but hardly recommended)
  • OpenCV with python bindings

Define a graph

This is how we define a network with a single convolution layer followed by ReLU activation function:

✔️

import tensorflow as tf
import os

# Input in NHWC (aka channels_last) data layout that means input blob dimensions
# are [batch x row x col x channel]. We defined input with 3 channels
# and 5x6 height and width correspondingly.
inp = tf.placeholder(tf.float32, [1, 5, 6, 3], 'input')

# Number of output channels is 4. Kernel size 3x3 with zero paddings to keep
# the 5x6 spatial sizes.
conv = tf.layers.conv2d(inp, filters=4, kernel_size=[3, 3],
                        activation=tf.nn.relu, padding='SAME',
                        bias_initializer=tf.random_normal_initializer())

with tf.Session() as sess:
    # Save a graph definition (once).
    tf.train.write_graph(sess.graph.as_graph_def(), "", "graph.pb")

    # Initialize all the weights.
    sess.run(tf.global_variables_initializer())

    # ...
    # Training process
    # ...

    # Save a checkpoint file that contains weights only. It can be done every
    # number of training iterations.
    saver = tf.train.Saver()
    saver.save(sess, os.path.join('.', 'model.ckpt'))

Freezing a graph

In the end of training we have a file graph.pb with a graph definition but without trained weights and a set of model* files that have actual weights. We need to merge them into a single binary file contains both graph definition and weights. Use freeze_graph tool from TensorFlow framework for it. All that you need is a name of output node. We can obtain it by printing certain node:

>>> print conv
Tensor("conv2d/Relu:0", shape=(1, 5, 6, 4), dtype=float32)
python ~/tensorflow/tensorflow/python/tools/freeze_graph.py \
    --input_graph graph.pb \
    --input_checkpoint model.ckpt \
    --output_graph frozen_graph.pb \
    --output_node_names "conv2d/Relu"

Deploy a graph in TensorFlow

Now we can share trained graph that can be used in the following way:

import tensorflow as tf import numpy as np

Read a graph from file.

with tf.gfile.FastGFile('frozen_graph.pb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read())

Set a graph into runtime.

tf.import_graph_def(graph_def, name='')

with tf.Session() as sess: # Generate input (just for an example). inpData = np.random.standard_normal([1, 5, 6, 3]).astype(np.float32)

# Get a references to input and output tensors.
inp = sess.graph.get_tensor_by_name('input:0')
out = sess.graph.get_tensor_by_name('conv2d/Relu:0')

# Run a graph up to the desired tensor.
outData = sess.run(out, feed_dict={inp: inpData})
# print outData.shape
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment