Add load_optim option to load checkpoint but not optimizer state
@@ -80,15 +80,17 @@ class Trainer(object):
self.lr_scheduler, self._num_updates, self._optim_history, extra_state,
)
- def load_checkpoint(self, filename):
+ def load_checkpoint(self, filename, load_optim=True):
"""Load all training state from a checkpoint file."""
- extra_state, self._optim_history, last_optim_state = \
+ extra_state, optim_history, last_optim_state = \
utils.load_model_state(filename, self.model)
if last_optim_state is not None:
# rebuild optimizer after loading model, since params may have changed
self._build_optimizer()
+ if load_optim:
+ self._optim_history = optim_history
# only reload optimizer and lr_scheduler if they match
last_optim = self._optim_history[-1]
if last_optim['criterion_name'] == self.criterion.__class__.__name__:
