[ModelZoo] Set Saver's parameter sharded=True in distributed training. (#954)

Signed-off-by: 泊霆 <hujunqi.hjq@alibaba-inc.com>

Author: Mesilenceki

modelzoo/bst/train.py

@@ -612,9 +612,10 @@ 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))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(

modelzoo/dbmtl/train.py

@@ -527,9 +527,10 @@ 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))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(

modelzoo/dcn/train.py

@@ -594,9 +594,10 @@ 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))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(

modelzoo/dcnv2/train.py

@@ -610,9 +610,10 @@ 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))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(

modelzoo/deepfm/train.py

@@ -472,9 +472,10 @@ 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))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(

modelzoo/dien/train.py

@@ -776,10 +776,10 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
+    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
-        local_init_op=tf.group(tf.tables_initializer(),
-                               tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max))
+        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))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(

modelzoo/din/train.py

@@ -594,10 +594,10 @@ def train(sess_config,
     hooks = []
     hooks.extend(input_hooks)
 
+    sharded_saver = tf_config != None
     scaffold = tf.train.Scaffold(
-        local_init_op=tf.group(tf.tables_initializer(),
-                               tf.local_variables_initializer(), data_init_op),
-        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max))
+        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))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(

modelzoo/dlrm/train.py

@@ -507,9 +507,10 @@ 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))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(

modelzoo/dssm/train.py

@@ -478,9 +478,10 @@ 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))
+        saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=sharded_saver))
 
     stop_hook = tf.train.StopAtStepHook(last_step=steps)
     log_hook = tf.train.LoggingTensorHook(

modelzoo/esmm/train.py

@@ -534,9 +534,10 @@ 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(), train_init_op),
-        saver=tf.train.Saver(max_to_keep=keep_checkpoint_max))
+        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))
 
     stop_hook = tf.train.StopAtStepHook(last_step=train_steps)
     log_hook = tf.train.LoggingTensorHook(