[Embedding] Check the sharded property of tf.train.Saver.

JackMoriarty · JackMoriarty · commit afdf80ef9952 · 2024-05-22T10:32:40.000+08:00
Signed-off-by: chenbangduo.cbd &lt;chenbangduo.cbd@alibaba-inc.com&gt;
diff --git a/modelzoo/bst/train.py b/modelzoo/bst/train.py
@@ -612,10 +612,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/dbmtl/train.py b/modelzoo/dbmtl/train.py
@@ -527,10 +527,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/dcn/train.py b/modelzoo/dcn/train.py
@@ -594,10 +594,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/dcnv2/train.py b/modelzoo/dcnv2/train.py
@@ -610,10 +610,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/deepfm/train.py b/modelzoo/deepfm/train.py
@@ -472,10 +472,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/dien/train.py b/modelzoo/dien/train.py
@@ -776,10 +776,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/din/train.py b/modelzoo/din/train.py
@@ -594,10 +594,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/dlrm/train.py b/modelzoo/dlrm/train.py
@@ -507,10 +507,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/dssm/train.py b/modelzoo/dssm/train.py
@@ -478,10 +478,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/esmm/train.py b/modelzoo/esmm/train.py
@@ -534,10 +534,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=train_steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/masknet/train.py b/modelzoo/masknet/train.py
@@ -529,10 +529,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/mlperf/train.py b/modelzoo/mlperf/train.py
@@ -522,10 +522,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/mmoe/train.py b/modelzoo/mmoe/train.py
@@ -523,10 +523,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/ple/train.py b/modelzoo/ple/train.py
@@ -592,10 +592,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/simple_multitask/train.py b/modelzoo/simple_multitask/train.py
@@ -427,10 +427,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=train_steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/modelzoo/wide_and_deep/train.py b/modelzoo/wide_and_deep/train.py
@@ -543,10 +543,9 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
-    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
         local_init_op=tf.group(tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(
diff --git a/tensorflow/python/training/saver.py b/tensorflow/python/training/saver.py
@@ -1071,10 +1071,14 @@ def _build(self, checkpoint_path, build_save, build_restore):
         # pylint: disable=protected-access
         self._var_list = variables._all_saveable_objects()
       from tensorflow.python.ops import hash_table
+      from tensorflow.python.ops import kv_variable_ops
       if isinstance(self._var_list, dict):
+        ev = {}
         ht = {}
         lst = {}
         for name, x in self._var_list.items():
+          if isinstance(x, kv_variable_ops.EmbeddingVariable):
+            ev[name] = x
           if isinstance(x, hash_table.HashTable):
             if x.hash_table not in ht:
               ht[x.hash_table] = [x]
@@ -1084,15 +1088,20 @@ def _build(self, checkpoint_path, build_save, build_restore):
             lst[name] = BloomFilterSaveable(x)
           else:
             lst[name] = x
+        if len(ev) != 0 and not self._sharded:
+          raise ValueError("EmbeddingVariable can only use sharded saver")
         if len(ht) != 0 and not self._sharded:
           raise ValueError("HashTable can only use sharded saver")
         for x, y in ht.items():
           lst[x.name] = HashTableSaveable(y)
         self._var_list = lst
       else:
+        ev = []
         ht = {}
         lst = []
         for x in self._var_list:
+          if isinstance(x, kv_variable_ops.EmbeddingVariable):
+            ev.append(x)
           if isinstance(x, hash_table.HashTable):
             if x.hash_table not in ht:
               ht[x.hash_table] = [x]
@@ -1102,6 +1111,8 @@ def _build(self, checkpoint_path, build_save, build_restore):
             lst.append(BloomFilterSaveable(x))
           else:
             lst.append(x)
+        if len(ev) != 0 and not self._sharded:
+          raise ValueError("EmbeddingVariable can only use sharded saver")
         if len(ht) != 0 and not self._sharded:
           raise ValueError("HashTable can only use sharded saver")
         for x, y in ht.items():
diff --git a/tensorflow/python/training/saver_test.py b/tensorflow/python/training/saver_test.py
@@ -852,6 +852,12 @@ def _model():
     for orig, restored in zip(orig_vals, restored_vals):
       self.assertAllEqual(orig, restored)
 
+  def testEnableSaverShardedWhenUseEmbeddingVariable(self):
+    with ops_lib.Graph().as_default():
+      emb_var = \
+        variable_scope.get_embedding_variable(name="emb_var", embedding_dim=64)
+      with self.assertRaisesRegexp(ValueError, "EmbeddingVariable"):
+        saver_module.Saver([emb_var], sharded=False)
 
 class SaveRestoreShardedTest(test.TestCase):
 
