def create_from_file(path):
try:
model_dict = load_torch_file(path)
except FileNotFoundError:
# If no model is found
raise FileNotFoundError(
'No valid model data found in {}'.format(path))
for model_name in Model.subclasses:
if model_name in model_dict:
model = Model.subclasses[model_name].from_dict(model_dict)
return model
def from_directory(cls, directory, device_id=None):
logger.info('Loading training state from {}'.format(directory))
root_path = Path(directory)
model_path = root_path / const.MODEL_FILE
model = Model.create_from_file(model_path)
if device_id is not None:
model.to(device_id)
optimizer_path = root_path / const.OPTIMIZER
optimizer_dict = load_torch_file(str(optimizer_path))
optimizer = optimizer_class(optimizer_dict['name'])(model.parameters(),
lr=0.0)
optimizer.load_state_dict(optimizer_dict['state_dict'])
trainer = cls(model, optimizer, checkpointer=None)
trainer_path = root_path / const.TRAINER
state = load_torch_file(str(trainer_path))
trainer.__dict__.update(state)
return trainer