add an utility to save pretrained model

Author: lucienwang1009

README.md

@@ -179,8 +179,9 @@ python tests/run_pretrained_models.py --backend onnxruntime --config tests/run_p
 ### <a name="save_pretrained_model"></a>Tool to save pre-trained model
 
 We provide an [utility](tools/save_pretrained_model.py) to save pre-trained model along with its config.
-Put `save_pretrained_model(sess, outputs, feed_inputs, save_dir, model_name)` in your last testing step and the pre-trained model and config will be saved under `save_dir/to_onnx`. 
-Please refer to [tools/save_pretrained_model.py](tools/save_pretrained_model.py) for more information.
+Put `save_pretrained_model(sess, outputs, feed_inputs, save_dir, model_name)` in your last testing epoch and the pre-trained model and config will be saved under `save_dir/to_onnx`. 
+Please refer to the example in [tools/save_pretrained_model.py](tools/save_pretrained_model.py) for more information.
+Note the minimum required Tensorflow version is r1.6.
 
 # Using the Python API
 ## TensorFlow to ONNX conversion

tools/save_pretrained_model.py

@@ -1,17 +1,46 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+"""
+Save pre-trained model.
+"""
 import tensorflow as tf
 import numpy as np
 
+# pylint: disable=redefined-outer-name,reimported
+
 def save_pretrained_model(sess, outputs, feeds, out_dir, model_name="pretrained"):
+    """Save pretrained model and config"""
     try:
         import os
+        import sys
+        import tensorflow as tf
+        import subprocess
         to_onnx_path = "{}/to_onnx".format(out_dir)
         if not os.path.isdir(to_onnx_path):
             os.makedirs(to_onnx_path)
         saved_model = "{}/saved_model".format(to_onnx_path)
         inputs_path = "{}/inputs.npy".format(to_onnx_path)
         pretrained_model_yaml_path = "{}/pretrained.yaml".format(to_onnx_path)
+        envars_path = "{}/environment.txt".format(to_onnx_path)
+        pip_requirement_path = "{}/requirements.txt".format(to_onnx_path)
+
+        print("===============Save Saved Model========================")
+        if os.path.exists(saved_model):
+            print("{} already exists, SKIP".format(saved_model))
+            return
+
+        print("Save tf version, python version and installed packages")
+        tf_version = tf.__version__
+        py_version = sys.version
+        pip_packages = subprocess.check_output([sys.executable, "-m", "pip", "freeze", "--all"])
+        pip_packages = pip_packages.decode("UTF-8")
+        with open(envars_path, "w") as fp:
+            fp.write(tf_version + os.linesep)
+            fp.write(py_version)
+        with open(pip_requirement_path, "w") as fp:
+            fp.write(pip_packages)
 
-        print("===============Save Frozen Graph========================")
         print("Save model for tf2onnx: {}".format(to_onnx_path))
         # save inputs
         inputs = {}
@@ -26,7 +55,7 @@ def save_pretrained_model(sess, outputs, feeds, out_dir, model_name="pretrained"
         # save graph and weights
         from tensorflow.saved_model import simple_save
         simple_save(sess, saved_model,
-                    {n: i for n,i in zip(inputs.keys(), feeds.keys())},
+                    {n: i for n, i in zip(inputs.keys(), feeds.keys())},
                     {op.name: op for op in outputs})
         print("Saved model to {}".format(saved_model))
 
@@ -39,9 +68,8 @@ def save_pretrained_model(sess, outputs, feeds, out_dir, model_name="pretrained"
 '''.format(model_name)
         pretrained_model_yaml += "  inputs:\n"
         for inp, _ in inputs.items():
-            pretrained_model_yaml += "    \"{}\": np.array(np.load(\"./inputs.npy\")[()][\"{}\"])\n".format(
-                inp, inp
-            )
+            pretrained_model_yaml += \
+                "    \"{input}\": np.array(np.load(\"./inputs.npy\")[()][\"{input}\"])\n".format(input=inp)
         outputs = [op.name for op in outputs]
         pretrained_model_yaml += "  outputs:\n"
         for out in outputs:
@@ -50,11 +78,12 @@ def save_pretrained_model(sess, outputs, feeds, out_dir, model_name="pretrained"
             f.write(pretrained_model_yaml)
         print("Saved pretrained model yaml to {}".format(pretrained_model_yaml_path))
         print("=========================================================")
-    except Exception as ex:
+    except Exception as ex:  # pylint: disable=broad-except
         print("Error: {}".format(ex))
 
 
 def test():
+    """Test sample."""
     x_val = np.random.rand(5, 20).astype(np.float32)
     y_val = np.random.rand(20, 10).astype(np.float32)
     x = tf.placeholder(tf.float32, x_val.shape, name="x")
@@ -66,8 +95,9 @@ def test():
     feeds = {x: x_val, y: y_val}
     with tf.Session() as sess:
         sess.run(init)
-        out = sess.run(outputs, feeds)
-        # NOTE: Put below snippet after the LAST testing step
+        sess.run(outputs, feeds)
+        # NOTE: NOT override the saved model, so put below snippet after testing the BEST model.
+        # if you perform testing several times.
         save_pretrained_model(sess, outputs, feeds, "./tests", model_name="test")