Merge pull request #27 from Yoctol/freeze_graph

we can freeze graph before saving

Author: SoluMilken

.travis.yml

@@ -13,7 +13,7 @@ before_install:
   - sudo apt-get install pandoc -y
 install:
   - pip install pipenv==2018.10.13
-  - make install
+  - make install-dev
   - pipenv run pip install tensorflow
 script:
   - pipenv run make lint

Makefile

@@ -2,6 +2,10 @@
 
 .PHONY: install
 install:
+	pipenv install
+
+.PHONY: install-dev
+install-dev:
 	pipenv install --dev
 
 .PHONY: lint
@@ -21,7 +25,7 @@ testall:
 	docker-compose rm -f
 
 .PHONY: all
-all: install testall lint
+all: install-dev lint testall
 
 .PHONY: clean
 clean:

README.md

@@ -13,15 +13,19 @@
 
 Some python utilities when using tensorflow-serving.
 
+
 ## Installation
 
+Prepare an environment with python version >= 3.6
+
 From PYPI:
 1. Manually install tensorflow CPU or GPU version.
 2. `pip install serving-utils`
 
 From Github repository:
 1. `git clone git@github.com:Yoctol/serving-utils.git`
-2. `make install`
+2. Manually install tensorflow CPU or GPU version.
+3. `make install`
 
 
 ## Usage
@@ -41,6 +45,7 @@ saver = Saver(
             outputs={'output': tf.Tensor...},
         )
     },
+    freeze=True,  # (default: True) Frozen graph will be saved if True.
 )
 saver.save(...)
 ```
@@ -60,6 +65,15 @@ client.predict(
 await client.async_predict(...)
 ```
 
+3. Freeze graph
+```python
+from serving_utils.freeze_graph import freeze_graph, create_session_from_graphdef
+
+frozen_graph_def = freeze_graph(session, output_op_names)
+new_session = create_session_from_graphdef(frozen_graph_def)
+```
+
+
 ## Test
 
 Run the following commands:
@@ -68,8 +82,9 @@ make lint
 make test
 ```
 
+
 ## Dev
 
 ```
-make install
+make install-dev
 ```

serving_utils/freeze_graph.py

@@ -0,0 +1,63 @@
+import tensorflow as tf
+from tensorflow.python.framework import graph_util
+
+
+def freeze_graph(session, output_op_names):
+    """Freeze graph
+
+    Freeze graph accroding to input session and
+    output operation names
+
+    ref: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/
+        python/tools/freeze_graph.py
+
+    Args:
+        session: an instance of tf.Session
+        output_op_names (strs): a list of operation names
+
+    Return:
+        a tf.GraphDef object
+
+    """
+    operations_in_graph(graph_def=session.graph_def, op_names=output_op_names)
+    frozen_graph_def = graph_util.convert_variables_to_constants(
+        sess=session,
+        input_graph_def=session.graph_def,
+        output_node_names=output_op_names,  # node name == op name
+    )
+    return frozen_graph_def
+
+
+def operations_in_graph(graph_def, op_names):
+    """
+    Check whether all operations are in graph or not
+
+    Args:
+        graph_def (tf.GraphDef): a tf.GraphDef object
+        op_names (strs): a list of operation names
+
+    Return:
+        KeyError or None
+
+    """
+    all_node_names = [node.name for node in graph_def.node]
+    for op_name in op_names:
+        if op_name not in all_node_names:
+            raise KeyError(f"{op_name} is not in this graph.")
+
+
+def create_session_from_graphdef(graph_def):
+    """
+    Create new session from given tf.GraphDef object
+
+    Arg:
+       graph_def (tf.GraphDef): a tf.GraphDef object
+
+    Return:
+       session (tf.Session): a new session with given graph_def
+
+    """
+    new_sess = tf.Session(graph=tf.Graph())
+    with new_sess.graph.as_default():
+        tf.import_graph_def(graph_def, name="")
+    return new_sess

serving_utils/saver.py

@@ -1,8 +1,9 @@
 from typing import Dict
 import pathlib
-
 import tensorflow as tf
 
+from .freeze_graph import freeze_graph, create_session_from_graphdef
+
 
 class Saver:
 
@@ -11,35 +12,71 @@ def __init__(
             session: tf.Session,
             output_dir: str,
             signature_def_map: Dict[str, tf.saved_model.signature_def_utils.predict_signature_def],
+            freeze: bool = True,
         ) -> None:
         self.session = session
         self.signature_def_map = signature_def_map
         self.output_dir = pathlib.Path(output_dir)
         self.output_dir.mkdir(parents=True, exist_ok=True)
+        self.freeze = freeze
 
-    @staticmethod
-    def _get_next_version(
-            path: pathlib.PosixPath,
-        ) -> pathlib.PosixPath:
-        candidate_paths = path.glob('**/*')
-        dirs = [x.name for x in candidate_paths if x.is_dir()]
-        versions = [int(dir_name) for dir_name in dirs if dir_name.isdigit()]
-        new_version = max(versions) + 1 if versions else 0
-        return path / str(new_version)
+    def save(self, **kwargs) -> str:
+        """Save model given session and signature
+
+        Save frozen model if self.freeze is True.
+
+        Return:
+            output_version_dir
+
+        """
+        if self.freeze:
+            output_op_names = self._get_op_names()
+            frozen_graphdef = freeze_graph(self.session, output_op_names)
+            session_to_be_saved = create_session_from_graphdef(frozen_graphdef)
+        else:
+            session_to_be_saved = self.session
 
-    def save(
-            self,
-            **kwargs
-        ) -> str:
         output_version_dir = str(self._get_next_version(self.output_dir))
-        with self.session.graph.as_default():
+        with session_to_be_saved.graph.as_default():
             builder = tf.saved_model.builder.SavedModelBuilder(
                 export_dir=output_version_dir,
             )
             builder.add_meta_graph_and_variables(
-                sess=self.session,
+                sess=session_to_be_saved,
                 tags=[tf.saved_model.tag_constants.SERVING],
                 signature_def_map=self.signature_def_map,
             )
         builder.save()
+
+        if self.freeze:
+            session_to_be_saved.close()
+
         return output_version_dir
+
+    @staticmethod
+    def _get_next_version(path: pathlib.PosixPath) -> pathlib.PosixPath:
+        candidate_paths = path.glob('**/*')
+        dirs = [x.name for x in candidate_paths if x.is_dir()]
+        versions = [int(dir_name) for dir_name in dirs if dir_name.isdigit()]
+        new_version = max(versions) + 1 if versions else 0
+        return path / str(new_version)
+
+    def _get_op_names(self):
+        """
+        Extract output operation names from signature
+
+        Return:
+            output_op_names (strs): a list of operation names
+                for output recorded in signature.
+
+        """
+        output_op_names = []
+        for signature_value in self.signature_def_map.values():
+            for _, tensor_info in signature_value.outputs.items():
+                tensor_name = tensor_info.name
+                if ':' in tensor_name:
+                    op_name = tensor_name.split(':')[0]
+                else:
+                    op_name = tensor_name
+                output_op_names.append(op_name)
+        return output_op_names

serving_utils/tests/mock_model.py

@@ -64,10 +64,9 @@ def fit(self, x, y, epochs):
                     self.y_place: y,
                 },
             )
-            self.save_tf_serve()
 
-    def save_tf_serve(self):
-        pass
+    def evaluate(self, x, y):
+        return self.sess.run(self.loss.name, feed_dict={self.x_place.name: x, self.y_place.name: y})
 
     def define_signature(self):
         signature_def_map = {

serving_utils/tests/test_client.py

@@ -0,0 +1,44 @@
+import pytest
+import numpy as np
+
+from ..freeze_graph import freeze_graph, create_session_from_graphdef
+from .mock_model import MockModel
+
+
+@pytest.fixture
+def x_test():
+    return np.random.rand(5, 10).astype('float32')
+
+
+@pytest.fixture
+def y_test(x_test):
+    batch_size, dims = x_test.shape
+    return np.random.rand(batch_size, dims).astype('float32')
+
+
+@pytest.fixture
+def model(x_test):
+    return MockModel(maxlen=x_test.shape[1])
+
+
+def test_create_session_from_graphdef(model):
+    old_graph_def = model.sess.graph_def
+    new_sess = create_session_from_graphdef(model.sess.graph_def)
+    assert old_graph_def == new_sess.graph_def
+
+
+def test_freeze_graph_will_not_change_loss(x_test, y_test, model):
+    old_loss = model.evaluate(x_test, y_test)
+    graph_def = freeze_graph(model.sess, [model.loss.op.name])
+
+    # model get new sess created from frozen graph def
+    new_sess = create_session_from_graphdef(graph_def)
+    model.sess = new_sess
+    new_loss = model.evaluate(x_test, y_test)
+
+    assert old_loss == new_loss
+
+
+def test_op_not_in_graph(model):
+    with pytest.raises(KeyError):
+        freeze_graph(model.sess, ['not_in_graph'])