Updating trainer so that train config parameter...

Updating trainer so that train config parameter `load_all_detection_checkpoint_vars` is used in restore_map. PiperOrigin-RevId: 187403903

Updating trainer so that train config parameter...
Updating trainer so that train config parameter `load_all_detection_checkpoint_vars` is used in restore_map. PiperOrigin-RevId: 187403903
65f20011 · Zhichao Lu · lzc5123016 · a9ee6b7e · 65f20011
Commit 65f20011 authored Feb 28, 2018 by Zhichao Lu Committed by lzc5123016 Mar 03, 2018
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

research/object_detection/trainer.py research/object_detection/trainer.py +3 -1

No files found.
--- a/research/object_detection/trainer.py
+++ b/research/object_detection/trainer.py
@@ -268,7 +268,9 @@ def train(create_tensor_dict_fn, create_model_fn, train_config, master, task,
    init_fn = None
    if train_config.fine_tune_checkpoint:
      var_map = detection_model.restore_map(
-          from_detection_checkpoint=train_config.from_detection_checkpoint)
+          from_detection_checkpoint=train_config.from_detection_checkpoint,
+          load_all_detection_checkpoint_vars=(
+              train_config.load_all_detection_checkpoint_vars))
      available_var_map = (variables_helper.
                           get_variables_available_in_checkpoint(
                               var_map, train_config.fine_tune_checkpoint))