added save_checkpoints_steps as flag (#747)

* added save_checkpoints_steps as flag

* added save_checkpoints_steps to TPUConfig

Author: sadransh

efficientdet/main.py

@@ -79,6 +79,8 @@
                      'evaluation.')
 flags.DEFINE_integer('iterations_per_loop', 100,
                      'Number of iterations per TPU training loop')
+flags.DEFINE_integer('save_checkpoints_steps', 5000,
+                     'Number of iterations per checkpoint save')
 flags.DEFINE_string(
     'training_file_pattern', None,
     'Glob for training data files (e.g., COCO train - minival set)')
@@ -218,6 +220,7 @@ def _can_partition(spatial_dim):
       FLAGS.iterations_per_loop if FLAGS.strategy == 'tpu' else 1,
       num_cores_per_replica=num_cores_per_replica,
       input_partition_dims=input_partition_dims,
+      save_checkpoints_steps=FLAGS.save_checkpoints_steps,
       per_host_input_for_training=tf.estimator.tpu.InputPipelineConfig
       .PER_HOST_V2)
 
@@ -240,7 +243,7 @@ def _can_partition(spatial_dim):
         train_distribute=strategy,
         log_step_count_steps=FLAGS.iterations_per_loop,
         session_config=config_proto,
-        save_checkpoints_steps=5000,
+        save_checkpoints_steps=FLAGS.save_checkpoints_steps,
         tf_random_seed=FLAGS.tf_random_seed,
     )