Example: Exporting models from Tensorpack (#885)

* start the showcase of exporting models from tensorpack

* mobile support

* showcase loading a frozen and pruned graph

* fix travis issues

* add docs

* add example image

* classes to functions

* typo

* typos

* some more fixes

* update docs

* rename the example

docs/tutorial/index.rst

@@ -22,6 +22,7 @@ User Tutorials
   save-load
   summary
   inference
+  export
   faq
 
 

docs/tutorial/inference.md

@@ -70,10 +70,59 @@ Alternatively, use tensorpack's `SaverRestore(path).init(tf.get_default_session(
 Tensorpack provides one tool [OfflinePredictor](../modules/predict.html#tensorpack.predict.OfflinePredictor),
 to merge the above two steps together.
 It has simple functionailities to build the graph, load the checkpoint, and return a callable for you.
-Check out examples and docs for its usage.
 
 OfflinePredictor is only for quick demo purposes.
 It runs inference on numpy arrays, therefore may not be the most efficient way.
 It also has very limited functionalities.
 If you need anything more complicated, please __do it on your own__ because Tensorpack
 doesn't care what happened after training.
+
+A simple explanation of how it works:
+```python
+pred_config = PredictConfig(
+    session_init=get_model_loader(model_path),
+    model=YourModel(),
+    input_names=['input1', 'input2'],
+    output_names=['output1', 'output2'])
+predictor = OfflinePredictor(pred_config)
+outputs = predictor(input1_array, input2_array)
+```
+
+As mentioned before, you might want to use a different graph for inference, 
+e.g., use NHWC format, support base64-encoded images. 
+You can make these changes in the `model` or `tower_func` in your `PredictConfig`.
+The example in [examples/basic/export-model.py](../examples/basic/export-model.py) demonstrates such an altered inference graph.
+
+### Exporter
+
+In addition to the standard checkpoint format tensorpack saved for you during training. 
+You can also save your models into other formats so it may be more friendly for inference.
+
+1. Export to `SavedModel` format for TensorFlow Serving:
+```python
+from tfutils.export import ModelExporter
+ModelExporter(pred_config).export_serving('/path/to/export')
+```
+
+This format contains both the graph and the variables. Refer to TensorFlow
+serving documentation on how to use it.
+
+2. Export to a frozen and pruned graph:
+
+```python
+ModelExporter(pred_config).export_compact('/path/to/compact_graph.pb')
+```
+
+This format is just a serialized `tf.Graph`. The export process:
+- Converts all variables to constants to embed the variables directly in the graph.
+- Removes all unnecessary operations (training-only ops, e.g., learning-rate) to compress the graph.
+
+This creates a self-contained graph which includes all necessary information to run inference.
+
+To load the graph, you can simply:
+```python
+graph_def = tf.GraphDef()
+graph_def.ParseFromString(open(graph_file, 'rb').read())
+tf.import_graph_def(graph_def)
+```
+[examples/basic/export-model.py](../examples/basic/export-model.py) demonstrates the usage of such a frozen/pruned graph.

examples/basics/export-model.py

+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import argparse
+import cv2
+import tensorflow as tf
+from tensorpack import *
+from tensorpack.tfutils.export import ModelExporter
+
+"""
+This example illustrates the process of exporting a model trained in Tensorpack to:
+- SavedModel format for TensorFlow Serving
+- A frozen and pruned inference graph (compact)
+
+The model applies a laplace filter to the input image.
+
+The steps are:
+
+1. train the model by
+
+    python export.py
+
+2. export the model by
+
+    python export.py --export serving --load train_log/export/checkpoint
+    python export.py --export compact --load train_log/export/checkpoint
+
+3. run inference by
+
+    python export.py --apply default --load train_log/export/checkpoint
+    python export.py --apply inference_graph --load train_log/export/checkpoint
+    python export.py --apply compact --load /tmp/compact_graph.pb
+"""
+
+
+SHAPE = 256
+CHANNELS = 3
+
+
+class Model(ModelDesc):
+    """Just a simple model, which applies the Laplacian-operation to images to showcase
+    the usage of variables, and alternating the inference-graph later.
+    """
+
+    def inputs(self):
+        return [tf.placeholder(tf.uint8, (None, SHAPE, SHAPE, CHANNELS), 'input_img'),
+                tf.placeholder(tf.uint8, (None, SHAPE, SHAPE, CHANNELS), 'target_img')]
+
+    def make_prediction(self, img):
+
+        img = tf.cast(img, tf.float32)
+        img = tf.image.rgb_to_grayscale(img)
+
+        k = tf.get_variable('filter', dtype=tf.float32,
+                            initializer=[[[[0.]], [[1.]], [[0.]]], [
+                                [[1.]], [[-4.]], [[1.]]], [[[0.]], [[1.]], [[0.]]]])
+        prediction_img = tf.nn.conv2d(img, k, strides=[1, 1, 1, 1], padding='SAME')
+        return prediction_img
+
+    def build_graph(self, input_img, target_img):
+
+        target_img = tf.cast(target_img, tf.float32)
+        target_img = tf.image.rgb_to_grayscale(target_img)
+
+        self.prediction_img = tf.identity(self.make_prediction(input_img), name='prediction_img')
+
+        cost = tf.losses.mean_squared_error(target_img, self.prediction_img,
+                                            reduction=tf.losses.Reduction.MEAN)
+        return tf.identity(cost, name='total_costs')
+
+    def optimizer(self):
+        lr = tf.get_variable('learning_rate', initializer=0.0, trainable=False)
+        return tf.train.AdamOptimizer(lr)
+
+
+def get_data(subset):
+    ds = FakeData([[SHAPE, SHAPE, CHANNELS], [SHAPE, SHAPE, CHANNELS]], 1000, random=False,
+                  dtype=['uint8', 'uint8'], domain=[(0, 255), (0, 10)])
+    ds = BatchData(ds, 1)
+    return ds
+
+
+class InferenceOnlyModel(Model):
+    """Recreate a different inference graph to accept images encoded as png. """
+
+    def inputs(self):
+        # The inference graph only accepts a single image, which is different to the training model.
+        return [tf.placeholder(tf.string, (None,), 'input_img_bytes')]
+
+    def build_graph(self, input_img_bytes):
+        # prepare input (png encoded strings to images)
+        input_img = tf.map_fn(lambda x: tf.image.decode_png(x, channels=3), input_img_bytes, dtype=tf.uint8)
+
+        # just copy the relevant parts to this graph.
+        prediction_img = self.make_prediction(input_img)
+
+        # outputs should be png encoded strings agains
+        prediction_img = tf.clip_by_value(prediction_img, 0, 255)
+        prediction_img = tf.cast(prediction_img, tf.uint8)
+        prediction_img_bytes = tf.map_fn(tf.image.encode_png, prediction_img, dtype=tf.string)
+
+        tf.identity(prediction_img_bytes, name='prediction_img_bytes')
+
+
+def export_serving(model_path):
+    """Export trained model to use it in TensorFlow Serving or cloudML. """
+    pred_config = PredictConfig(
+        session_init=get_model_loader(model_path),
+        model=InferenceOnlyModel(),
+        input_names=['input_img_bytes'],
+        output_names=['prediction_img_bytes'])
+    ModelExporter(pred_config).export_serving('/tmp/exported')
+
+
+def export_compact(model_path):
+    """Export trained model to use it as a frozen and pruned inference graph in
+       mobile applications. """
+    pred_config = PredictConfig(
+        session_init=get_model_loader(model_path),
+        model=Model(),
+        input_names=['input_img'],
+        output_names=['prediction_img'])
+    ModelExporter(pred_config).export_compact('/tmp/compact_graph.pb')
+
+
+def apply(model_path):
+    """Run inference from a training model checkpoint. """
+    pred_config = PredictConfig(
+        session_init=get_model_loader(model_path),
+        model=Model(),
+        input_names=['input_img'],
+        output_names=['prediction_img'])
+
+    pred = OfflinePredictor(pred_config)
+    img = cv2.imread('lena.png')
+    prediction = pred([img])[0]
+    cv2.imwrite('applied_default.jpg', prediction[0])
+
+
+def apply_inference_graph(model_path):
+    """Run inference from a different graph, which receives encoded images buffers. """
+    pred_config = PredictConfig(
+        session_init=get_model_loader(model_path),
+        model=InferenceOnlyModel(),
+        input_names=['input_img_bytes'],
+        output_names=['prediction_img_bytes'])
+
+    pred = OfflinePredictor(pred_config)
+    buf = open('lena.png', 'rb').read()
+    prediction = pred([buf])[0]
+    with open('applied_inference_graph.png', 'wb') as f:
+        f.write(prediction[0])
+
+
+def apply_compact(graph_path):
+    """Run the pruned and frozen inference graph. """
+    with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
+        # Note, we just load the graph and do *not* need to initialize anything.
+        with tf.gfile.GFile(graph_path, "rb") as f:
+            graph_def = tf.GraphDef()
+            graph_def.ParseFromString(f.read())
+            tf.import_graph_def(graph_def)
+
+        input_img = sess.graph.get_tensor_by_name('import/input_img:0')
+        prediction_img = sess.graph.get_tensor_by_name('import/prediction_img:0')
+
+        prediction = sess.run(prediction_img, {input_img: cv2.imread('lena.png')[None, ...]})
+        cv2.imwrite('applied_compact.png', prediction[0])
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--load', help='load model')
+    parser.add_argument('--apply', help='run sampling', default='',
+                        choices=['default', 'inference_graph', 'compact'])
+    parser.add_argument('--export', help='export the model', default='',
+                        choices=['serving', 'compact'])
+
+    args = parser.parse_args()
+
+    if args.apply != '':
+        if args.apply == 'default':
+            apply(args.load)
+        elif args.apply == 'inference_graph':
+            apply_inference_graph(args.load)
+        else:
+            apply_compact(args.load)
+    elif args.export != '':
+        if args.export == 'serving':
+            export_serving(args.load)
+        else:
+            export_compact(args.load)
+    else:
+        logger.auto_set_dir()
+
+        ds_train = get_data('train')
+
+        config = TrainConfig(
+            model=Model(),
+            data=QueueInput(ds_train),
+            callbacks=[
+                ModelSaver(),
+            ],
+            steps_per_epoch=1,
+            max_epoch=1,
+        )
+        launch_train_with_config(config, SimpleTrainer())

tensorpack/tfutils/export.py

 # -*- coding: utf-8 -*-
 # File: export.py
-# Author: Patrick Wieschollek <mail@patwie.com>
 
 """
-This simplifies the process of exporting a model for TensorFlow serving.
+A collection of functions to ease the process of exporting
+a model for production.
 
 """
 
 import tensorflow as tf
+from tensorflow.python.framework import graph_util
+from tensorflow.python.platform import gfile
+from tensorflow.python.tools import optimize_for_inference_lib
+
 from ..utils import logger
-from ..graph_builder.model_desc import ModelDescBase
+from ..tfutils.common import get_tensors_by_names
+from ..tfutils.tower import PredictTowerContext
 from ..input_source import PlaceholderInput
-from ..tfutils import TowerContext, sessinit
 
+__all__ = ['ModelExporter']
 
-__all__ = ['ModelExport']
 
+class ModelExporter(object):
+    """Export models for inference."""
 
-class ModelExport(object):
-    """Wrapper for tf.saved_model"""
-    def __init__(self, model, input_names, output_names):
+    def __init__(self, config):
         """Initialise the export process.
 
-        Example:
-
-            .. code-block:: python
-                from mnist_superresolution import Model
-                from tensorpack.tfutils import export
-
-                e = ModelExport(Model(), ['lowres'], ['prediction'])
-                e.export('train_log/mnist_superresolution/checkpoint', 'export/first_export')
-
-            Will generate a model for TensorFlow serving with input 'lowres' and
-            output 'prediction'. The model is in the directory 'export' and can be
-            loaded by
-
-            .. code-block:: python
-
-                import tensorflow as tf
-                from tensorflow.python.saved_model import tag_constants
-
-                export_dir = 'export/first_export'
-                with tf.Session(graph=tf.Graph(), config=tf.ConfigProto(allow_soft_placement=True)) as sess:
-                    tf.saved_model.loader.load(sess, [tag_constants.SERVING], export_dir)
-
-                    prediction = tf.get_default_graph().get_tensor_by_name('prediction:0')
-                    lowres = tf.get_default_graph().get_tensor_by_name('lowres:0')
-
-                    prediction = sess.run(prediction, {lowres: ...})[0]
-
         Args:
-            model (ModelDescBase): the model description which should be exported
-            input_names (list(str)): names of input tensors
-            output_names (list(str)): names of output tensors
+            config (PredictConfig): the config to use.
+                The graph will be built with `config.tower_func` and `config.inputs_desc`.
+                Then the input / output names will be used to export models for inference.
         """
+        super(ModelExporter, self).__init__()
+        self.config = config
 
-        assert isinstance(input_names, list)
-        assert isinstance(output_names, list)
-        assert isinstance(model, ModelDescBase)
-        self.model = model
-        self.output_names = output_names
-        self.input_names = input_names
+    def export_compact(self, filename):
+        """Create a self-contained inference-only graph and write final graph (in pb format) to disk.
 
-    def export(self, checkpoint, export_path,
+        Args:
+            filename (str): path to the output graph
+        """
+        self.graph = self.config._maybe_create_graph()
+        with self.graph.as_default():
+            input = PlaceholderInput()
+            input.setup(self.config.inputs_desc)
+            with PredictTowerContext(''):
+                self.config.tower_func(*input.get_input_tensors())
+
+            input_tensors = get_tensors_by_names(self.config.input_names)
+            output_tensors = get_tensors_by_names(self.config.output_names)
+
+            self.config.session_init._setup_graph()
+            # we cannot use "self.config.session_creator.create_session()" here since it finalizes the graph
+            sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
+            self.config.session_init._run_init(sess)
+
+            dtypes = [n.dtype for n in input_tensors]
+
+            # freeze variables to constants
+            frozen_graph_def = graph_util.convert_variables_to_constants(
+                sess,
+                self.graph.as_graph_def(),
+                [n.name[:-2] for n in output_tensors],
+                variable_names_whitelist=None,
+                variable_names_blacklist=None)
+
+            # prune unused nodes from graph
+            pruned_graph_def = optimize_for_inference_lib.optimize_for_inference(
+                frozen_graph_def,
+                [n.name[:-2] for n in input_tensors],
+                [n.name[:-2] for n in output_tensors],
+                [dtype.as_datatype_enum for dtype in dtypes],
+                False)
+
+            with gfile.FastGFile(filename, "wb") as f:
+                f.write(pruned_graph_def.SerializeToString())
+                logger.info("Output graph written to {}.".format(filename))
+
+    def export_serving(self, filename,
                        tags=[tf.saved_model.tag_constants.SERVING],
                        signature_name='prediction_pipeline'):
         """
+        Converts a checkpoint and graph to a servable for TensorFlow Serving.
         Use SavedModelBuilder to export a trained model without tensorpack dependency.
 
         Remarks:
@@ -79,56 +95,37 @@ class ModelExport(object):
             https://github.com/tensorflow/serving/blob/master/tensorflow_serving/g3doc/signature_defs.md
 
         Args:
-            checkpoint (str): path to checkpoint file
-            export_path (str): path for export directory
+            filename (str): path for export directory
             tags (list): list of user specified tags
             signature_name (str): name of signature for prediction
         """
-        logger.info('[export] build model for %s' % checkpoint)
-        with TowerContext('', is_training=False):
+
+        self.graph = self.config._maybe_create_graph()
+        with self.graph.as_default():
             input = PlaceholderInput()
-            input.setup(self.model.get_inputs_desc())
-            self.model.build_graph(*input.get_input_tensors())
-
-        self.sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
-        # load values from latest checkpoint
-        init = sessinit.SaverRestore(checkpoint)
-        self.sess.run(tf.global_variables_initializer())
-        init.init(self.sess)
-
-        self.inputs = []
-        for n in self.input_names:
-            tensor = tf.get_default_graph().get_tensor_by_name('%s:0' % n)
-            logger.info('[export] add input-tensor "%s"' % tensor.name)
-            self.inputs.append(tensor)
-
-        self.outputs = []
-        for n in self.output_names:
-            tensor = tf.get_default_graph().get_tensor_by_name('%s:0' % n)
-            logger.info('[export] add output-tensor "%s"' % tensor.name)
-            self.outputs.append(tensor)
-
-        logger.info('[export] exporting trained model to %s' % export_path)
-        builder = tf.saved_model.builder.SavedModelBuilder(export_path)
-
-        logger.info('[export] build signatures')
-        # build inputs
-        inputs_signature = dict()
-        for n, v in zip(self.input_names, self.inputs):
-            logger.info('[export] add input signature: %s' % v)
-            inputs_signature[n] = tf.saved_model.utils.build_tensor_info(v)
-
-        outputs_signature = dict()
-        for n, v in zip(self.output_names, self.outputs):
-            logger.info('[export] add output signature: %s' % v)
-            outputs_signature[n] = tf.saved_model.utils.build_tensor_info(v)
+            input.setup(self.config.inputs_desc)
+            with PredictTowerContext(''):
+                self.config.tower_func(*input.get_input_tensors())
+
+            input_tensors = get_tensors_by_names(self.config.input_names)
+            inputs_signatures = {t.name: tf.saved_model.utils.build_tensor_info(t) for t in input_tensors}
+            output_tensors = get_tensors_by_names(self.config.output_names)
+            outputs_signatures = {t.name: tf.saved_model.utils.build_tensor_info(t) for t in output_tensors}
+
+            self.config.session_init._setup_graph()
+            # we cannot use "self.config.session_creator.create_session()" here since it finalizes the graph
+            sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
+            self.config.session_init._run_init(sess)
+
+            builder = tf.saved_model.builder.SavedModelBuilder(filename)
 
             prediction_signature = tf.saved_model.signature_def_utils.build_signature_def(
-            inputs=inputs_signature,
-            outputs=outputs_signature,
+                inputs=inputs_signatures,
+                outputs=outputs_signatures,
                 method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)
 
             builder.add_meta_graph_and_variables(
-            self.sess, tags,
+                sess, tags,
                 signature_def_map={signature_name: prediction_signature})
             builder.save()
+            logger.info("SavedModel created at {}.".format(filename))