[Embedding] Check the sharded property of tf.train.Saver. (#996)

Signed-off-by: chenbangduo.cbd <[email protected]>

Author: JackMoriarty

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(

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(

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(

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(

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(

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(

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(

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(

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(

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(

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(

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(

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(

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(

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(

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(

tensorflow/python/feature_column/feature_column_v2_test.py

@@ -7527,7 +7527,7 @@ def testEmbeddingVariableForL2FeatureEviction(self):
     opt = ftrl.FtrlOptimizer(0.1, l1_regularization_strength=2.0, l2_regularization_strength=0.00001)
     g_v = opt.compute_gradients(loss)
     train_op = opt.apply_gradients(g_v)
-    saver = saver_module.Saver()
+    saver = saver_module.Saver(sharded=True)
     init = variables_lib.global_variables_initializer()
     with self.test_session() as sess:
       sess.run(ops.get_collection(ops.GraphKeys.EV_INIT_VAR_OPS))
@@ -7758,7 +7758,7 @@ def testEmbeddingVariableForSharedEmbeddingColumnsWithPartitionNum(self):
     g_v = opt.compute_gradients(loss)
     train_op = opt.apply_gradients(g_v)
     init = variables_lib.global_variables_initializer()
-    saver = saver_module.Saver()
+    saver = saver_module.Saver(sharded=True)
 
   @test_util.run_deprecated_v1
   def testEmbeddingVariableForInt32ID(self):
@@ -7783,7 +7783,7 @@ def testEmbeddingVariableForInt32ID(self):
     opt = ftrl.FtrlOptimizer(0.1, l1_regularization_strength=2.0, l2_regularization_strength=0.00001)
     g_v = opt.compute_gradients(loss)
     train_op = opt.apply_gradients(g_v)
-    saver = saver_module.Saver()
+    saver = saver_module.Saver(sharded=True)
     init = variables_lib.global_variables_initializer()
     with self.test_session() as sess:
       sess.run(ops.get_collection(ops.GraphKeys.EV_INIT_VAR_OPS))

tensorflow/python/ops/embedding_variable_ops_gpu_test.py

@@ -63,7 +63,8 @@ def testEmbeddingVariableForInitFromProto(self):
     g_v = opt.compute_gradients(loss)
     train_op = opt.apply_gradients(g_v)
     graph = ops.get_default_graph()
-    meta_graph_def = saver_module.export_meta_graph()
+    saver = saver_module.Saver(sharded=True)
+    meta_graph_def = saver_module.export_meta_graph(saver_def=saver.as_saver_def())
     ops.reset_default_graph()
     with self.test_session() as sess:
       res = saver_module.import_meta_graph(meta_graph_def)
@@ -748,7 +749,7 @@ def testSaveV3(self):
     g_v = opt.compute_gradients(loss)
     train_op = opt.apply_gradients(g_v, global_step=gs)
     init = variables.global_variables_initializer()
-    saver = saver = saver_module.Saver()
+    saver = saver = saver_module.Saver(sharded=True)
     checkpoint_directory = self.get_temp_dir()
     model_path = os.path.join(checkpoint_directory, "model.ckpt")
     with self.test_session() as sess:
@@ -816,7 +817,7 @@ def testEmbeddingVariableSaveAndRestoreOptimzierStatesForMultiTierWithHbm(self):
         opt = adagrad.AdagradOptimizer(0.1)
         g_v = opt.compute_gradients(loss)
         train_op = opt.apply_gradients(g_v, gs)
-        saver = saver_module.Saver()
+        saver = saver_module.Saver(sharded=True)
         graph = ops.get_default_graph()
         with self.test_session(graph = graph) as sess:
           saver.restore(sess, os.path.join(checkpoint_directory, "model.ckpt-12345"))