unify temp test data saving logic

Author: nbcsm

tests/backend_test_base.py

@@ -40,6 +40,14 @@ def setUp(self):
             tf.logging.set_verbosity(tf.logging.WARN)
             self.log.setLevel(logging.INFO)
 
+    def tearDown(self):
+        if not self.config.is_debug_mode:
+            utils.delete_directory(self.test_data_directory)
+
+    @property
+    def test_data_directory(self):
+        return os.path.join(self.config.temp_dir, self._testMethodName)
+
     @staticmethod
     def assertAllClose(expected, actual, **kwargs):
         np.testing.assert_allclose(expected, actual, **kwargs)
@@ -72,6 +80,7 @@ def run_onnxruntime(self, model_path, inputs, output_names):
     def _run_backend(self, g, outputs, input_dict):
         model_proto = g.make_model("test")
         model_path = self.save_onnx_model(model_proto, input_dict)
+
         if self.config.backend == "onnxmsrtnext":
             y = self.run_onnxmsrtnext(model_path, input_dict, outputs)
         elif self.config.backend == "onnxruntime":
@@ -93,8 +102,6 @@ def run_test_case(self, feed_dict, input_names_with_port, output_names_with_port
             onnx_feed_dict = feed_dict
 
         graph_def = None
-        save_dir = os.path.join(self.config.temp_path, self._testMethodName)
-
         if convert_var_to_const:
             with tf.Session() as sess:
                 variables_lib.global_variables_initializer().run()
@@ -113,20 +120,18 @@ def run_test_case(self, feed_dict, input_names_with_port, output_names_with_port
             expected = sess.run(output_dict, feed_dict=feed_dict)
 
         if self.config.is_debug_mode:
-            if not os.path.exists(save_dir):
-                os.makedirs(save_dir)
-            model_path = os.path.join(save_dir, self._testMethodName + "_original.pb")
-            with open(model_path, "wb") as f:
-                f.write(sess.graph_def.SerializeToString())
+            if not os.path.exists(self.test_data_directory):
+                os.makedirs(self.test_data_directory)
+            model_path = os.path.join(self.test_data_directory, self._testMethodName + "_original.pb")
+            utils.save_protobuf(model_path, sess.graph_def)
             self.log.debug("created file %s", model_path)
 
         graph_def = tf_optimize(input_names_with_port, output_names_with_port,
                                 sess.graph_def, constant_fold)
 
         if self.config.is_debug_mode and constant_fold:
-            model_path = os.path.join(save_dir, self._testMethodName + "_after_tf_optimize.pb")
-            with open(model_path, "wb") as f:
-                f.write(graph_def.SerializeToString())
+            model_path = os.path.join(self.test_data_directory, self._testMethodName + "_after_tf_optimize.pb")
+            utils.save_protobuf(model_path, graph_def)
             self.log.debug("created file  %s", model_path)
 
         tf.reset_default_graph()
@@ -146,9 +151,8 @@ def run_test_case(self, feed_dict, input_names_with_port, output_names_with_port
                 self.assertEqual(expected_val.shape, actual_val.shape)
 
     def save_onnx_model(self, model_proto, feed_dict, postfix=""):
-        save_path = os.path.join(self.config.temp_path, self._testMethodName)
-        target_path = utils.save_onnx_model(save_path, self._testMethodName + postfix, feed_dict, model_proto,
-                                            include_test_data=self.config.is_debug_mode,
+        target_path = utils.save_onnx_model(self.test_data_directory, self._testMethodName + postfix, feed_dict,
+                                            model_proto, include_test_data=self.config.is_debug_mode,
                                             as_text=self.config.is_debug_mode)
 
         self.log.debug("create model file: %s", target_path)

tests/common.py

@@ -6,15 +6,16 @@
 import argparse
 import os
 import sys
-import tempfile
 import unittest
 
 from distutils.version import LooseVersion
+from tf2onnx import utils
 from tf2onnx.tfonnx import DEFAULT_TARGET, POSSIBLE_TARGETS
 
 __all__ = ["TestConfig", "get_test_config", "unittest_main",
-           "check_tf_min_version", "check_opset_min_version", "check_target", "skip_onnxruntime_backend",
-           "skip_caffe2_backend", "check_onnxruntime_incompatibility"]
+           "check_tf_min_version", "skip_tf_versions",
+           "check_opset_min_version", "check_target", "skip_onnxruntime_backend", "skip_caffe2_backend",
+           "check_onnxruntime_incompatibility"]
 
 
 # pylint: disable=missing-docstring
@@ -28,7 +29,7 @@ def __init__(self):
         self.backend = os.environ.get("TF2ONNX_TEST_BACKEND", "onnxruntime")
         self.backend_version = self._get_backend_version()
         self.is_debug_mode = False
-        self.temp_path = tempfile.mkdtemp()
+        self.temp_dir = utils.get_temp_directory()
 
     @property
     def is_mac(self):
@@ -67,28 +68,32 @@ def __str__(self):
                             "target={}".format(self.target),
                             "backend={}".format(self.backend),
                             "backend_version={}".format(self.backend_version),
-                            "is_debug_mode={}".format(self.is_debug_mode)])
+                            "is_debug_mode={}".format(self.is_debug_mode),
+                            "temp_dir={}".format(self.temp_dir)])
 
     @staticmethod
     def load():
         config = TestConfig()
         # if not launched by pytest, parse console arguments to override config
         if "pytest" not in sys.argv[0]:
             parser = argparse.ArgumentParser()
-            parser.add_argument('--backend', default=config.backend,
+            parser.add_argument("--backend", default=config.backend,
                                 choices=["caffe2", "onnxmsrtnext", "onnxruntime"],
                                 help="backend to test against")
-            parser.add_argument('--opset', type=int, default=config.opset, help="opset to test against")
+            parser.add_argument("--opset", type=int, default=config.opset, help="opset to test against")
             parser.add_argument("--target", default=",".join(config.target), choices=POSSIBLE_TARGETS,
                                 help="target platform")
             parser.add_argument("--debug", help="output debugging information", action="store_true")
-            parser.add_argument('unittest_args', nargs='*')
+            parser.add_argument("--temp_dir", help="temp dir")
+            parser.add_argument("unittest_args", nargs='*')
 
             args = parser.parse_args()
             config.backend = args.backend
             config.opset = args.opset
             config.target = args.target.split(',')
             config.is_debug_mode = args.debug
+            if args.temp_dir:
+                config.temp_dir = args.temp_dir
 
             # Now set the sys.argv to the unittest_args (leaving sys.argv[0] alone)
             sys.argv[1:] = args.unittest_args

tests/run_pretrained_models.py

@@ -10,10 +10,8 @@
 
 import argparse
 import os
-import shutil
 import sys
 import tarfile
-import tempfile
 import time
 import traceback
 import zipfile
@@ -37,7 +35,7 @@
 
 # pylint: disable=broad-except,logging-not-lazy,unused-argument,unnecessary-lambda
 
-TMPPATH = tempfile.mkdtemp()
+TEMP_DIR = os.path.join(utils.get_temp_directory(), "run_pretrained")
 PERFITER = 1000
 
 
@@ -192,7 +190,7 @@ def run_caffe2(self, name, model_proto, inputs):
     def run_onnxmsrtnext(self, name, model_proto, inputs):
         """Run test against msrt-next backend."""
         import lotus
-        model_path = utils.save_onnx_model(TMPPATH, name, inputs, model_proto)
+        model_path = utils.save_onnx_model(TEMP_DIR, name, inputs, model_proto)
         m = lotus.InferenceSession(model_path)
         results = m.run(self.output_names, inputs)
         if self.perf:
@@ -205,8 +203,8 @@ def run_onnxmsrtnext(self, name, model_proto, inputs):
     def run_onnxruntime(self, name, model_proto, inputs):
         """Run test against msrt-next backend."""
         import onnxruntime as rt
-        model_path = utils.save_onnx_model(TMPPATH, name, inputs, model_proto, include_test_data=True)
-        utils.save_onnx_model(TMPPATH, name, inputs, model_proto, include_test_data=False, as_text=True)
+        model_path = utils.save_onnx_model(TEMP_DIR, name, inputs, model_proto, include_test_data=True)
+        utils.save_onnx_model(TEMP_DIR, name, inputs, model_proto, include_test_data=False, as_text=True)
         print("\t\t" + model_path)
         m = rt.InferenceSession(model_path)
         results = m.run(self.output_names, inputs)
@@ -221,8 +219,7 @@ def run_onnxruntime(self, name, model_proto, inputs):
     def create_onnx_file(name, model_proto, inputs, outdir):
         os.makedirs(outdir, exist_ok=True)
         model_path = os.path.join(outdir, name + ".onnx")
-        with open(model_path, "wb") as f:
-            f.write(model_proto.SerializeToString())
+        utils.save_protobuf(model_path, model_proto)
         print("\tcreated", model_path)
 
     def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=None, perf=None, fold_const=None):
@@ -440,8 +437,8 @@ def main():
             ret = None
             print(ex)
         finally:
-            if os.path.exists(TMPPATH) and not args.debug:
-                shutil.rmtree(TMPPATH)
+            if not args.debug:
+                utils.delete_directory(TEMP_DIR)
         if not ret:
             failed += 1
 

tests/test_backend.py

@@ -722,8 +722,8 @@ def test_relu(self):
     @check_onnxruntime_incompatibility("Mul")
     def test_leaky_relu(self):
         for alpha in [0.1, -0.1, 1.0, -1.0, 10.0, -10.0]:
-            x_val = 1000*np.random.random_sample([1000, 100]).astype(np.float32)
-            x = tf.placeholder(tf.float32, [None]*x_val.ndim, name=_TFINPUT)
+            x_val = 1000 * np.random.random_sample([1000, 100]).astype(np.float32)
+            x = tf.placeholder(tf.float32, [None] * x_val.ndim, name=_TFINPUT)
             x_ = tf.nn.leaky_relu(x, alpha)
             _ = tf.identity(x_, name=_TFOUTPUT)
             self._run_test_case([_OUTPUT], {_INPUT: x_val})
@@ -1199,7 +1199,7 @@ def test_strided_slice6(self):
 
     @skip_caffe2_backend("multiple dims not supported")
     def test_strided_slice7(self):
-        x_val = np.arange(5*6).astype("float32").reshape(5, 6)
+        x_val = np.arange(5 * 6).astype("float32").reshape(5, 6)
 
         x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
         x_ = tf.strided_slice(x, [0, 1], [3, 4], [1, 1], begin_mask=2)
@@ -1218,7 +1218,6 @@ def test_strided_slice7(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-
     @skip_caffe2_backend("fails with schema error")
     @check_opset_min_version(7, "batchnorm")
     def test_batchnorm(self):

tf2onnx/convert.py

@@ -14,12 +14,14 @@
 from onnx import helper
 import tensorflow as tf
 
-import tf2onnx.utils
+import tf2onnx
+from tf2onnx import utils
 from tf2onnx.graph import GraphUtil
 from tf2onnx.tfonnx import process_tf_graph, tf_optimize, DEFAULT_TARGET, POSSIBLE_TARGETS
 
 _TENSORFLOW_DOMAIN = "ai.onnx.converters.tensorflow"
 
+
 # pylint: disable=unused-argument
 
 
@@ -45,7 +47,7 @@ def get_args():
 
     args.shape_override = None
     if args.inputs:
-        args.inputs, args.shape_override = tf2onnx.utils.split_nodename_and_shape(args.inputs)
+        args.inputs, args.shape_override = utils.split_nodename_and_shape(args.inputs)
     if args.outputs:
         args.outputs = args.outputs.split(",")
     if args.inputs_as_nchw:
@@ -64,14 +66,14 @@ def default_custom_op_handler(ctx, node, name, args):
 def main():
     args = get_args()
 
-    opset = tf2onnx.utils.find_opset(args.opset)
+    opset = utils.find_opset(args.opset)
     print("using tensorflow={}, onnx={}, opset={}, tfonnx={}/{}".format(
         tf.__version__, onnx.__version__, opset,
         tf2onnx.__version__, tf2onnx.version.git_version[:6]))
 
     # override unknown dimensions from -1 to 1 (aka batchsize 1) since not every runtime does
     # support unknown dimensions.
-    tf2onnx.utils.ONNX_UNKNOWN_DIMENSION = args.unknown_dim
+    utils.ONNX_UNKNOWN_DIMENSION = args.unknown_dim
 
     if args.custom_ops:
         # default custom ops for tensorflow-onnx are in the "tf" namespace
@@ -113,9 +115,8 @@ def main():
 
     # write onnx graph
     if args.output:
-        with open(args.output, "wb") as f:
-            f.write(model_proto.SerializeToString())
-            print("\nComplete successfully, the onnx model is generated at " + args.output)
+        utils.save_protobuf(args.output, model_proto)
+        print("\nComplete successfully, the onnx model is generated at " + args.output)
 
 
 if __name__ == "__main__":

tf2onnx/utils.py

@@ -11,6 +11,8 @@
 
 import os
 import re
+import shutil
+import tempfile
 import six
 import numpy as np
 import tensorflow as tf
@@ -95,6 +97,7 @@
 # Fake onnx op type which is used for Graph input.
 GRAPH_INPUT_TYPE = "NON_EXISTENT_ONNX_TYPE"
 
+
 def make_name(name):
     """Make op name for inserted ops."""
     global INTERNAL_NAME
@@ -274,18 +277,13 @@ def save_onnx_model(save_path_root, onnx_file_name, feed_dict, model_proto, incl
             t = numpy_helper.from_array(data)
             t.name = data_key
             data_full_path = os.path.join(data_path, "input_" + str(i) + ".pb")
-            with open(data_full_path, 'wb') as f:
-                f.write(t.SerializeToString())
+            save_protobuf(data_full_path, t)
             i += 1
 
     target_path = os.path.join(save_path, onnx_file_name + ".onnx")
-    with open(target_path, "wb") as f:
-        f.write(model_proto.SerializeToString())
-
+    save_protobuf(target_path, model_proto)
     if as_text:
-        with open(target_path + ".pbtxt", "w") as f:
-            f.write(text_format.MessageToString(model_proto))
-
+        save_protobuf(target_path + ".pbtxt", model_proto, as_text=True)
     return target_path
 
 
@@ -347,3 +345,22 @@ def tf_name_scope(name):
 def create_vague_shape_like(shape):
     make_sure(len(shape) >= 0, "rank should be >= 0")
     return [-1 for i in enumerate(shape)]
+
+
+def get_temp_directory():
+    return os.environ.get("TF2ONNX_TEMP_DIRECTORY", tempfile.mkdtemp())
+
+
+def delete_directory(path):
+    if os.path.exists(path):
+        shutil.rmtree(path)
+
+
+def save_protobuf(path, message, as_text=False):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    if as_text:
+        with open(path, "w") as f:
+            f.write(text_format.MessageToString(message))
+    else:
+        with open(path, "wb") as f:
+            f.write(message.SerializeToString())