Fix loading older fp16 optimizer checkpoints.
@@ -668,21 +668,43 @@ class TorchAgent(ABC, Agent):
if self.fp16:
self.optimizer = fp16_optimizer_wrapper(self.optimizer)
- if optim_states:
- if saved_optim_type != opt['optimizer']:
- print('WARNING: not loading optim state since optim class '
- 'changed.')
+ # TODO: we might want to hard reset optimizers here in the
+ # case of fine tuning. Some rudimentary experiments seemed to
+ # indicate that keeping adam weights around was desirable, so this
+ # will remain the behavior for the time being.
+ if optim_states and saved_optim_type != opt['optimizer']:
+ # we changed from adam to adamax, or sgd to adam, or similar
+ print('WARNING: not loading optim state since optim class changed.')
+ elif optim_states:
+ # check for any fp16/fp32 conversions we need to do
+ optimstate_fp16 = ('loss_scaler' in optim_states)
+ if self.fp16 and optimstate_fp16:
+ # previously trained in fp16, now we're training in fp16.
+ # ideally no action needed, but APEX broke backwards
+ # compatibility and this is the hack around it.
+ optim_states['loss_scaler'] = (
+ self.optimizer.state_dict()['loss_scaler']
+ )
+ elif optimstate_fp16 and not self.fp16:
+ # old optimizer was fp16 but now we're doing fp32,
+ # drop the fp16 wrapper from the state_dict and just load
+ # the fp16 weights into the fp32 tensors
+ optim_states = optim_states['optimizer_state_dict']
+ elif not optimstate_fp16 and self.fp16:
+ # old optimizer was fp32, but now we're doing fp16.
+ # this is a bit clunky, but alternatives are worse
+ self.optimizer.optimizer.load_state_dict(optim_states)
+ return
else:
+ # previously trained in fp32, loading in fp32.
+ # no special treatment needed.
+ pass
+
+ # finally, try to actually load the optimizer state
try:
self.optimizer.load_state_dict(optim_states)
except ValueError:
- print('WARNING: not loading optim state since model '
- 'params changed.')
- if self.use_cuda:
- for state in self.optimizer.state.values():
- for k, v in state.items():
- if isinstance(v, torch.Tensor):
- state[k] = v.cuda()
+ print('WARNING: not loading optim state since model params changed.')
def build_lr_scheduler(self, states=None, hard_reset=False):
"""
@@ -1377,7 +1399,8 @@ class TorchAgent(ABC, Agent):
else:
states['model'] = self.model.state_dict()
- if hasattr(self, 'optimizer'): # save optimizer params
+ if hasattr(self, 'optimizer'):
+ # save optimizer params
states['optimizer'] = self.optimizer.state_dict()
states['optimizer_type'] = self.opt['optimizer']
