add test case test_table_save_load_local_file_system_for_estimator

Author: PWZER

tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_variable_test.py

@@ -58,6 +58,8 @@
 from tensorflow.python.training.tracking import data_structures
 from tensorflow.python.training.tracking import util as track_util
 from tensorflow.python.util import compat
+from tensorflow_estimator.python.estimator import estimator
+from tensorflow_estimator.python.estimator import estimator_lib
 
 try:
   import tensorflow_io
@@ -958,6 +960,53 @@ def test_table_save_load_local_file_system(self):
 
       del table
 
+  def test_table_save_load_local_file_system_for_estimator(self):
+
+    def input_fn():
+      return {"x": constant_op.constant([1], dtype=dtypes.int64)}
+
+    def model_fn(features, labels, mode, params):
+      file_system_saver = de.FileSystemSaver()
+      embedding = de.get_variable(
+          name="embedding",
+          dim=3,
+          trainable=False,
+          key_dtype=dtypes.int64,
+          value_dtype=dtypes.float32,
+          initializer=-1.0,
+          kv_creator=de.CuckooHashTableCreator(saver=file_system_saver),
+      )
+      lookup = de.embedding_lookup(embedding, features["x"])
+      upsert = embedding.upsert(features["x"],
+                                constant_op.constant([[1.0, 2.0, 3.0]]))
+
+      with ops.control_dependencies([lookup, upsert]):
+        train_op = state_ops.assign_add(training.get_global_step(), 1)
+
+      scaffold = training.Scaffold(
+          saver=saver.Saver(sharded=True,
+                            max_to_keep=1,
+                            keep_checkpoint_every_n_hours=None,
+                            defer_build=True,
+                            save_relative_paths=True))
+      est = estimator_lib.EstimatorSpec(mode=mode,
+                                        scaffold=scaffold,
+                                        loss=constant_op.constant(0.),
+                                        train_op=train_op,
+                                        predictions=lookup)
+      return est
+
+    save_dir = os.path.join(self.get_temp_dir(), "save_restore")
+    save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash")
+
+    # train and save
+    est = estimator.Estimator(model_fn=model_fn, model_dir=save_path)
+    est.train(input_fn=input_fn, steps=1)
+
+    # restore and predict
+    predict_results = next(est.predict(input_fn=input_fn))
+    self.assertAllEqual(predict_results, [1.0, 2.0, 3.0])
+
   def test_save_restore_only_table(self):
     if context.executing_eagerly():
       self.skipTest('skip eager test when using legacy Saver.')

tensorflow_recommenders_addons/dynamic_embedding/python/ops/tf_save_restore_patch.py

@@ -179,23 +179,26 @@ def _get_dynamic_embedding_save_ops(self):
 
     for var in self._var_list:
       de_var = None
-      if isinstance(var, (de.FileSystemSaver._DynamicEmbeddingShardFileSystemSaveable,
-                          de.FileSystemSaver._DynamicEmbeddingVariabelFileSystemSaveable)):
+      if isinstance(
+          var,
+          (de.FileSystemSaver._DynamicEmbeddingShardFileSystemSaveable,
+           de.FileSystemSaver._DynamicEmbeddingVariabelFileSystemSaveable)):
         de_var = var._de_variable
       elif isinstance(var, de.Variable) and var._saveable_object_creator:
         de_var = var
 
-      if de_var and isinstance(de_var._saveable_object_creator, de.FileSystemSaver):
+      if de_var and isinstance(de_var._saveable_object_creator,
+                               de.FileSystemSaver):
         if de_var._saveable_object_creator.config.save_path:
           de_variable_folder_dir = de_var._saveable_object_creator.config.save_path
         else:
           de_variable_folder_dir = self._de_var_fs_save_dir
 
         save_op = de_var.save_to_file_system(
-          dirpath=de_variable_folder_dir,
-          proc_size=de_var._saveable_object_creator.config.proc_size,
-          proc_rank=de_var._saveable_object_creator.config.proc_rank,
-          buffer_size=de_var._saveable_object_creator.config.buffer_size)
+            dirpath=de_variable_folder_dir,
+            proc_size=de_var._saveable_object_creator.config.proc_size,
+            proc_rank=de_var._saveable_object_creator.config.proc_rank,
+            buffer_size=de_var._saveable_object_creator.config.buffer_size)
         save_ops.as_list().append(save_op)
     return control_flow_ops.group(save_ops.as_list())
 
@@ -206,33 +209,36 @@ def _get_dynamic_embedding_restore_ops(self):
 
     for var in self._var_list:
       de_var = None
-      if isinstance(var, (de.FileSystemSaver._DynamicEmbeddingShardFileSystemSaveable,
-                          de.FileSystemSaver._DynamicEmbeddingVariabelFileSystemSaveable)):
+      if isinstance(
+          var,
+          (de.FileSystemSaver._DynamicEmbeddingShardFileSystemSaveable,
+           de.FileSystemSaver._DynamicEmbeddingVariabelFileSystemSaveable)):
         de_var = var._de_variable
       elif isinstance(var, de.Variable) and var._saveable_object_creator:
         de_var = var
 
-      if de_var and isinstance(de_var._saveable_object_creator, de.FileSystemSaver):
+      if de_var and isinstance(de_var._saveable_object_creator,
+                               de.FileSystemSaver):
         if de_var._saveable_object_creator.config.save_path:
           de_variable_folder_dir = de_var._saveable_object_creator.config.save_path
         else:
           de_variable_folder_dir = self._de_var_fs_save_dir
 
         restore_op = de_var.load_from_file_system_with_restore_function(
-          dirpath=de_variable_folder_dir,
-          proc_size=de_var._saveable_object_creator.config.proc_size,
-          proc_rank=de_var._saveable_object_creator.config.proc_rank,
-          buffer_size=de_var._saveable_object_creator.config.buffer_size)
+            dirpath=de_variable_folder_dir,
+            proc_size=de_var._saveable_object_creator.config.proc_size,
+            proc_rank=de_var._saveable_object_creator.config.proc_rank,
+            buffer_size=de_var._saveable_object_creator.config.buffer_size)
         restore_ops.as_list().append(restore_op)
     return control_flow_ops.group(restore_ops.as_list())
 
   def _build(self, checkpoint_path, build_save, build_restore):
-    super(_DynamicEmbeddingSaver, self)._build(
-      checkpoint_path, build_save, build_restore)
+    super(_DynamicEmbeddingSaver, self)._build(checkpoint_path, build_save,
+                                               build_restore)
 
     with ops.name_scope("FileSystemSaver", "save_to_file_system", []) as name:
       self._de_var_fs_save_dir = array_ops.placeholder(
-        dtype=dtypes.string, shape=(), name="de_var_file_system_save_dir")
+          dtype=dtypes.string, shape=(), name="de_var_file_system_save_dir")
       self._de_save_ops = self._get_dynamic_embedding_save_ops()
       self._de_restore_ops = self._get_dynamic_embedding_restore_ops()
 
@@ -337,14 +343,14 @@ def save(self,
 
     if global_step is not None:
       de_variable_folder_dir = os.path.join(
-        save_path_parent, "TFRADynamicEmbedding-{}".format(global_step))
+          save_path_parent, "TFRADynamicEmbedding-{}".format(global_step))
       if self._pad_step_number:
         # Zero-pads the step numbers, so that they are sorted when listed.
         de_variable_folder_dir = os.path.join(
-          save_path_parent, "TFRADynamicEmbedding-{:08d}".format(global_step))
+            save_path_parent, "TFRADynamicEmbedding-{:08d}".format(global_step))
     else:
-      de_variable_folder_dir = os.path.join(
-        save_path_parent, "TFRADynamicEmbedding")
+      de_variable_folder_dir = os.path.join(save_path_parent,
+                                            "TFRADynamicEmbedding")
 
     if not self._is_empty:
       try: