new option to replace optimizer

Former-commit-id: 5aabb5a584b873cf86abae21fa21d9d9c10c3f73
javiribera · Dec 5, 2018 · 9f2be09 · 9f2be09
1 parent 3fd4ca4
commit 9f2be09
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/object-locator/argparser.py b/object-locator/argparser.py
@@ -135,6 +135,12 @@ def parse_command_args(training_or_testing):
                             metavar='OPTIM',
                             choices=['sgd', 'adam'],
                             help='SGD or Adam')
+        parser.add_argument('--replace-optimizer',
+                            action='store_true',
+                            default=False,
+                            help='Replace optimizer state when resuming from checkpoint. '
+                                 'If True, the optimizer will be replaced as provided in the '
+                                 'arguments of this scripts. If not resuming, it has no effect.')
         parser.add_argument('--max-mask-pts',
                             type=int,
                             default=np.infty,

diff --git a/object-locator/train.py b/object-locator/train.py
@@ -146,7 +146,8 @@
                 lowest_mahd = np.infty
                 print('W: Loaded checkpoint has not been validated. ', end='')
             model.load_state_dict(checkpoint['model'])
-            optimizer.load_state_dict(checkpoint['optimizer'])
+            if not args.replace_optimizer:
+                optimizer.load_state_dict(checkpoint['optimizer'])
             print(f"\n\__ loaded checkpoint '{args.resume}'"
                   f"(now on epoch {checkpoint['epoch']})")
         else: