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