[Tensorflow] ] Replace tensors of a SavedModel or FrozenGraph

When working with pretrained deep learning models in TensorFlow the input tensor is often fixed and changing a tensor from the beginning of the graph is by design painful. In this post I want to show how to replace tensors of a graph and build a servable model for TensorFlow serving.

Getting a pretrained model

Here is a selection of pretrained object detection models TensorFlow. As the most of them are trained in TensorFlow 1.x and not working in TensorFlow 2, I will be focusing on TensorFlow 1.x. Some might also work in TensorFlow 2. For this example, I used ssd_mobilenet_v1_ppn_coco, which is a fast object detection model trained on the coco dataset.

Loading the model

First, we load the model. We have two options here, loading the frozen graph for inference or loading the SavedModel. Note that you can not re-train a frozen graph because variables are constants, but you can train a SavedModel. It depends on the objective and sometimes you only have one and not both available. In the end we need a GraphDef Object to proceed, how we get it, doesn’t matter.

import tensorflow as tf
sess = tf.Session()

def load_graph(frozen_graph_filename):
    with tf.gfile.GFile(frozen_graph_filename, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name='')
    return graph
graph_model = load_graph("ssd_mobilenet_v1_ppn\\frozen_inference_graph.pb")
graph_model_def = graph_model.as_graph_def()

With SavedModel:

model = tf.saved_model.load(sess, export_dir="ssd_mobilenet_v1_ppn/saved_model", tags=['serve'])
graph_model_def = graph_model.as_graph_def()

Create a new Graph

Now we create a new graph and add our tensors. In this example I will add base64/jpeg decoding and image resizing to make the model capable of accepting base64 encoded jpegs directly.

graph_base64 = tf.Graph()
graph_base64.as_default()

string_inp = tf.placeholder(tf.string, shape=(None,), name='base64_in')
imgs_map = tf.map_fn(
    tf.image.decode_image,
    string_inp,
    dtype=tf.uint8
)
imgs_map.set_shape((None, None, None, 3))
imgs = tf.image.resize_images(imgs_map, [300, 300], method=tf.image.ResizeMethod.BILINEAR)
imgs = tf.reshape(imgs, (-1, 300, 300, 3))
img_uint8 = tf.image.convert_image_dtype(imgs, dtype=tf.uint8, saturate=False)

Replace tensors

Now we have two graphs one contains the pretrained model and one our new input pipeline. To combine them we need to import the pretrained model into our new graph, while connecting the nodes. This is quite a way to go until it works properly, but here is how to achieve it:

tf.import_graph_def(graph_model_def, name='', input_map={"image_tensor": img_uint8})

What `import_graph_def` does is to load the pretrained graph into our new graph, while replacing the image_tensor with our input pipeline (`img_unit8` tensor). Keeping the name parameter empty is important to avoid duplication of nodes. It’s also possible to connect the output of the pretrained model with additional output tensors with (‘return_elements’), which is not necessary here.

Export for TensorFlow X (Serving)

Finally, we export the model as a servable SavedModel, which can be deployed in Tensorflow Serve.

builder = tf.saved_model.builder.SavedModelBuilder('coco_obj_det')
num_detections_tensor = sess.graph.get_tensor_by_name('num_detections:0')
detection_boxes_tensor = sess.graph.get_tensor_by_name('detection_boxes:0')
detection_scores_tensor = sess.graph.get_tensor_by_name('detection_scores:0')
detection_classes_tensor = sess.graph.get_tensor_by_name('detection_classes:0')
signature = tf.saved_model.signature_def_utils.predict_signature_def(                                                                        
    inputs={'base64_in': string_inp}, outputs={'num_detections': num_detections_tensor, "detection_boxes" : detection_boxes_tensor,
                                               'detection_scores' : detection_scores_tensor, 'detection_classes' : detection_classes_tensor})

builder.add_meta_graph_and_variables(                                                                                                        
    sess=sess,                                                                                                                    
    tags=[tf.saved_model.tag_constants.SERVING],                                                                                             
    signature_def_map={                                                                                                                      
        tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:                                                                
            signature                                                                                                                        
    })                                                                                                                                       
builder.save()

Be careful with your input pipeline when replacing tensors

Pretrained models often require a special preprocessing of inputs. For this example, the replacement of the input pipeline works, but it might not work with other pretrained models, without proper preprocessing of the data. For example using the same approach for the ssd_resnet_50_fpn_coco model won’t work, as the input pipeline is different. To get it working, it’s necessary to include the hole input pipeline from the pipeline.config file into the graph to replace the tensors.