set states in load function
@@ -122,31 +122,29 @@ class LanguageModelAgent(Agent):
print('[ Using CUDA ]')
torch.cuda.set_device(opt['gpu'])
+ init_model = None
# check first for 'init_model' for loading model from file
if opt.get('init_model') and os.path.isfile(opt['init_model']):
init_model = opt['init_model']
- # next check for 'model_file'
- elif opt.get('model_file') and os.path.isfile(opt['model_file']):
+ # next check for 'model_file', this would override init_model
+ if opt.get('model_file') and os.path.isfile(opt['model_file']):
init_model = opt['model_file']
- else:
- init_model = None
- if init_model is not None:
+ if init_model is not None and \
+ not os.path.isfile(init_model + '.opt'):
+ # this will only be called for loading older models before
+ # .opt file is saved
+ new_opt = self.load_opt(init_model)
# load model parameters if available
- print('Loading existing model params from ' + init_model)
- new_opt, self.states = self.load(init_model)
- # override model-specific options with stored ones if not
- # already overriden with .opt file
- if not os.path.isfile(init_model + '.opt'):
+ print('[ Setting opt from {} ]'.format(
+ init_model
+ ))
opt = self.override_opt(new_opt)
- if opt['dict_file'] is None:
- if init_model is not None and os.path.isfile(init_model + '.dict'):
- # check first to see if a dictionary exists
+ if (init_model is not None and
+ os.path.isfile(init_model + '.dict')) or \
+ opt['dict_file'] is None:
opt['dict_file'] = init_model + '.dict'
- elif opt.get('model_file'):
- # otherwise, set default dict-file if it is not set
- opt['dict_file'] = opt['model_file'] + '.dict'
# load dictionary and basic tokens & vectors
self.dict = DictionaryAgent(opt)
@@ -164,9 +162,8 @@ class LanguageModelAgent(Agent):
# set model
self.model = RNNModel(opt, len(self.dict))
- if self.states:
- # set loaded states if applicable
- self.model.load_state_dict(self.states['model'])
+ if init_model is not None:
+ self.load(init_model)
if self.use_cuda:
self.model.cuda()
@@ -581,7 +578,18 @@ class LanguageModelAgent(Agent):
if 'loss' in metrics_dict and self.scheduler is not None:
self.scheduler.step(metrics_dict['loss'])
- def load(self, path):
- """Return opt and model states."""
+ def load_opt(self, path):
+ """Return opt, states."""
states = torch.load(path, map_location=lambda cpu, _: cpu)
- return states['opt'], states
+ return states['opt']
+
+ def load(self, path):
+ """Load model states."""
+ if os.path.isfile(path):
+ # load model parameters if available
+ print('[ Loading existing model params from {} ]'.format(path))
+ self.states = torch.load(path, map_location=lambda cpu, _: cpu)
+
+ if self.states:
+ # set loaded states if applicable
+ self.model.load_state_dict(self.states['model'])
