diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index f2cb517f..e9224e2b 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -81,10 +81,10 @@ class TrainerIO(object): # request what to save from the model model = self.__get_model() checkpoint_dict = model.get_save_dict() - - # merge trainer and model saving items - checkpoint['state_dict'] = checkpoint_dict checkpoint.update(checkpoint_dict) + + # add the state_dict from the model + checkpoint['state_dict'] = checkpoint_dict return checkpoint # -------------------- @@ -186,6 +186,7 @@ class TrainerIO(object): # call model hook model.on_hpc_load() + def max_ckpt_in_folder(self, path): files = os.listdir(path) files = [x for x in files if 'ckpt_' in x] diff --git a/tests/test_models.py b/tests/test_models.py index e5781d32..1a315625 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -43,6 +43,7 @@ def test_loading_meta_tags(): clear_save_dir() + def test_dp_output_reduce(): # test identity when we have a single gpu @@ -64,6 +65,61 @@ def test_dp_output_reduce(): assert reduced['b']['c'] == out['b']['c'] +def test_model_saving_loading(): + hparams = get_hparams() + model = LightningTestModel(hparams) + + save_dir = init_save_dir() + + # exp file to get meta + exp = get_exp(False) + exp.argparse(hparams) + exp.save() + + trainer_options = dict( + max_nb_epochs=1, + cluster=SlurmCluster(), + experiment=exp, + checkpoint_callback=ModelCheckpoint(save_dir) + ) + + # fit model + trainer = Trainer(**trainer_options) + result = trainer.fit(model) + real_global_step = trainer.global_step + + # traning complete + assert result == 1, 'amp + ddp model failed to complete' + + # make a prediction + for batch in model.test_dataloader: + break + + x, y = batch + x = x.view(x.size(0), -1) + + # generate preds before saving model + model.eval() + pred_before_saving = model(x) + + # save model + new_weights_path = os.path.join(save_dir, 'save_test.ckpt') + trainer.save_checkpoint(new_weights_path) + + # load new model + tags_path = exp.get_data_path(exp.name, exp.version) + tags_path = os.path.join(tags_path, 'meta_tags.csv') + model_2 = LightningTestModel.load_from_metrics(weights_path=new_weights_path, tags_csv=tags_path, on_gpu=False) + model_2.eval() + + # make prediction + # assert that both predictions are the same + new_pred = model_2(x) + assert torch.eq(pred_before_saving, new_pred) + + clear_save_dir() + + def test_cpu_slurm_saving_loading(): """ Verify model save/load/checkpoint on CPU