diff --git a/tfjs-converter/python/run-python-tests.sh b/tfjs-converter/python/run-python-tests.sh index dfd9f79534a..7d1fd39c5c9 100755 --- a/tfjs-converter/python/run-python-tests.sh +++ b/tfjs-converter/python/run-python-tests.sh @@ -24,7 +24,7 @@ TEST_FILES="$(find "${SCRIPTS_DIR}" -name '*_test.py')" pip install virtualenv -TMP_VENV_DIR="$(mktemp -d --suffix=_venv)" +TMP_VENV_DIR="$(mktemp -u).venv" virtualenv -p "python" "${TMP_VENV_DIR}" source "${TMP_VENV_DIR}/bin/activate" diff --git a/tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py b/tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py index 6111944799f..4370a0f22f0 100644 --- a/tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py +++ b/tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2.py @@ -29,6 +29,8 @@ from tensorflow.python.grappler import cluster as gcluster from tensorflow.python.grappler import tf_optimizer from tensorflow.python.saved_model.load import load +from tensorflow.python.saved_model import loader +from tensorflow.python.tools import saved_model_utils from tensorflow.python.training.saver import export_meta_graph from google.protobuf.json_format import MessageToDict import tensorflow_hub as hub @@ -272,15 +274,20 @@ def _check_signature_in_model(saved_model, signature_name): saved_model.signatures.keys())) -def _freeze_saved_model_v1(graph, output_node_names): - frozen_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants( - tf.compat.v1.Session(), graph.as_graph_def(), output_node_names) +def _freeze_saved_model_v1(saved_model_dir, saved_model_tags, + output_node_names): + with tf.compat.v1.Session() as sess: + loader.load(sess, saved_model_tags, saved_model_dir) + input_graph_def = saved_model_utils.get_meta_graph_def( + saved_model_dir, ','.join(saved_model_tags)).graph_def + frozen_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants( + sess, input_graph_def, output_node_names) - frozen_graph = tf.Graph() - with frozen_graph.as_default(): - tf.import_graph_def(frozen_graph_def, name='') + frozen_graph = tf.Graph() + with frozen_graph.as_default(): + tf.import_graph_def(frozen_graph_def, name='') - return frozen_graph + return frozen_graph def _freeze_saved_model_v2(concrete_func): return convert_to_constants.convert_variables_to_constants_v2( @@ -336,8 +343,8 @@ def convert_tf_saved_model(saved_model_dir, try: frozen_graph = _freeze_saved_model_v2(concrete_func) except BaseException: - frozen_graph = _freeze_saved_model_v1( - concrete_func.graph, output_node_names) + frozen_graph = _freeze_saved_model_v1(saved_model_dir, saved_model_tags, + output_node_names) optimize_graph(frozen_graph, output_node_names, output_graph, model.tensorflow_version, diff --git a/tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2_test.py b/tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2_test.py index 5144cfb87c0..21edc83388c 100644 --- a/tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2_test.py +++ b/tfjs-converter/python/tensorflowjs/converters/tf_saved_model_conversion_v2_test.py @@ -80,6 +80,44 @@ def _create_saved_model_v1(self): builder.save() + def _create_saved_model_v1_with_hashtable(self): + """Create a TensorFlow SavedModel V1 with unused hash table for testing.""" + + graph = tf.Graph() + with graph.as_default(): + x = tf.placeholder('float32', [2, 2]) + w = tf.compat.v1.get_variable('w', shape=[2, 2]) + output = tf.compat.v1.matmul(x, w) + init_op = w.initializer + + # Add a hash table that is not used by the output. + keys = tf.constant(['key']) + values = tf.constant([1]) + initializer = tf.contrib.lookup.KeyValueTensorInitializer(keys, values) + table = tf.contrib.lookup.HashTable(initializer, -1) + + # Create a builder. + save_dir = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) + builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(save_dir) + + with tf.compat.v1.Session() as sess: + # Run the initializer on `w`. + sess.run(init_op) + table.init.run() + + builder.add_meta_graph_and_variables( + sess, [tf.compat.v1.saved_model.tag_constants.SERVING], + signature_def_map={ + "serving_default": + tf.compat.v1.saved_model \ + .signature_def_utils.predict_signature_def( + inputs={"x": x}, + outputs={"output": output}) + }, + assets_collection=None) + + builder.save() + def _create_saved_model_with_fusable_conv2d(self): """Test a basic model with fusable conv2d.""" layers = [ @@ -192,22 +230,24 @@ def double_module_fn(): def test_convert_saved_model_v1(self): self._create_saved_model_v1() + input_dir = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) + output_dir = os.path.join(input_dir, 'js') tf_saved_model_conversion_v2.convert_tf_saved_model( - os.path.join(self._tmp_dir, SAVED_MODEL_DIR), - os.path.join(self._tmp_dir, SAVED_MODEL_DIR) + input_dir, + output_dir ) - weights = [{ + expected_weights_manifest = [{ 'paths': ['group1-shard1of1.bin'], 'weights': [{'dtype': 'float32', 'name': 'w', 'shape': [2, 2]}]}] - tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) + tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR, 'js') # Check model.json and weights manifest. with open(os.path.join(tfjs_path, 'model.json'), 'rt') as f: model_json = json.load(f) self.assertTrue(model_json['modelTopology']) weights_manifest = model_json['weightsManifest'] - self.assertEqual(weights_manifest, weights) + self.assertEqual(weights_manifest, expected_weights_manifest) # Check meta-data in the artifact JSON. self.assertEqual(model_json['format'], 'graph-model') self.assertEqual( @@ -215,9 +255,37 @@ def test_convert_saved_model_v1(self): 'TensorFlow.js Converter v%s' % version.version) self.assertEqual(model_json['generatedBy'], tf.__version__) - self.assertTrue( - glob.glob( - os.path.join(self._tmp_dir, SAVED_MODEL_DIR, 'group*-*'))) + self.assertTrue(glob.glob(os.path.join(output_dir, 'group*-*'))) + + def test_convert_saved_model_v1_with_hashtable(self): + self._create_saved_model_v1_with_hashtable() + + input_dir = os.path.join(self._tmp_dir, SAVED_MODEL_DIR) + output_dir = os.path.join(input_dir, 'js') + tf_saved_model_conversion_v2.convert_tf_saved_model( + input_dir, + output_dir + ) + + expected_weights_manifest = [{ + 'paths': ['group1-shard1of1.bin'], + 'weights': [{'dtype': 'float32', 'name': 'w', 'shape': [2, 2]}]}] + + tfjs_path = os.path.join(self._tmp_dir, SAVED_MODEL_DIR, 'js') + # Check model.json and weights manifest. + with open(os.path.join(tfjs_path, 'model.json'), 'rt') as f: + model_json = json.load(f) + self.assertTrue(model_json['modelTopology']) + weights_manifest = model_json['weightsManifest'] + self.assertEqual(weights_manifest, expected_weights_manifest) + # Check meta-data in the artifact JSON. + self.assertEqual(model_json['format'], 'graph-model') + self.assertEqual( + model_json['convertedBy'], + 'TensorFlow.js Converter v%s' % version.version) + self.assertEqual(model_json['generatedBy'], + tf.__version__) + self.assertTrue(glob.glob(os.path.join(output_dir, 'group*-*'))) def test_convert_saved_model(self): self._create_saved_model()