From e2c7fa44b716bf5ff77f9d042a0fef66577478ee Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 21:37:06 -0400 Subject: [PATCH 01/24] auto state-dict and remove the way the model is loaded during hpc --- pytorch_lightning/root_module/model_saving.py | 14 +++++++++++--- pytorch_lightning/root_module/root_module.py | 7 +++++-- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index 25377851..ec55d5f6 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -9,17 +9,19 @@ class ModelIO(object): def load_model_specific(self, checkpoint): """ Do something with the checkpoint + Gives model a chance to load something before state_dict is restored :param checkpoint: :return: """ - raise NotImplementedError + pass def get_save_dict(self): """ Return specific things for the model + Called before trainer requests the state_dict :return: """ - raise NotImplementedError + pass # ------------------------- # OPTIONAL HOOKS @@ -80,6 +82,7 @@ class TrainerIO(object): checkpoint_dict = model.get_save_dict() # merge trainer and model saving items + checkpoint['state_dict'] = checkpoint_dict checkpoint.update(checkpoint_dict) return checkpoint @@ -167,13 +170,18 @@ class TrainerIO(object): else: checkpoint = torch.load(filepath, map_location=lambda storage, loc: storage) - # load training state + # load training state (affects trainer only) self.restore_training_state(checkpoint) # load model state model = self.__get_model() + + # give model a chance to load something model.load_model_specific(checkpoint) + # load the state_dict on the model automatically + model.load_state_dict(checkpoint['state_dict']) + # call model hook model.on_hpc_load() diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py index d6d740f0..a25f7f6f 100644 --- a/pytorch_lightning/root_module/root_module.py +++ b/pytorch_lightning/root_module/root_module.py @@ -110,9 +110,12 @@ class LightningModule(GradInformation, ModelIO, ModelHooks): model = cls(hparams) - # allow model to load + # give model a chance to load something model.load_model_specific(checkpoint) - model.load_state_dict(checkpoint['state_dict'], strict=False) + + # load the state_dict on the model automatically + model.load_state_dict(checkpoint['state_dict']) + return model def summarize(self): From aacf1947ea056fa53122faa8dab2fd5843da7044 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 21:38:06 -0400 Subject: [PATCH 02/24] auto state-dict and remove the way the model is loaded during hpc --- pytorch_lightning/root_module/model_saving.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index ec55d5f6..f2cb517f 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -4,6 +4,7 @@ import re import pdb from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel + class ModelIO(object): def load_model_specific(self, checkpoint): From 92a1f559b5101a6ca1ab62dc957c41489373f49a Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 21:39:01 -0400 Subject: [PATCH 03/24] remove state_dict --- pytorch_lightning/testing_models/lm_test_module.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/pytorch_lightning/testing_models/lm_test_module.py b/pytorch_lightning/testing_models/lm_test_module.py index e33ee53e..39fcbb4c 100644 --- a/pytorch_lightning/testing_models/lm_test_module.py +++ b/pytorch_lightning/testing_models/lm_test_module.py @@ -171,17 +171,6 @@ class LightningTestModel(LightningModule): def on_tng_metrics(self, logs): logs['some_tensor_to_test'] = torch.rand(1) - # --------------------- - # MODEL SAVING - # --------------------- - def get_save_dict(self): - checkpoint = {'state_dict': self.state_dict()} - return checkpoint - - def load_model_specific(self, checkpoint): - self.load_state_dict(checkpoint['state_dict']) - pass - # --------------------- # TRAINING SETUP # --------------------- From a5a80f35ec2c1474300ac08a9444b69213367da7 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 21:39:28 -0400 Subject: [PATCH 04/24] removed old template --- .../models/sample_model_template/__init__.py | 0 .../sample_model_template/model_template.py | 203 ------------------ 2 files changed, 203 deletions(-) delete mode 100644 pytorch_lightning/models/sample_model_template/__init__.py delete mode 100644 pytorch_lightning/models/sample_model_template/model_template.py diff --git a/pytorch_lightning/models/sample_model_template/__init__.py b/pytorch_lightning/models/sample_model_template/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/pytorch_lightning/models/sample_model_template/model_template.py b/pytorch_lightning/models/sample_model_template/model_template.py deleted file mode 100644 index 10f12c59..00000000 --- a/pytorch_lightning/models/sample_model_template/model_template.py +++ /dev/null @@ -1,203 +0,0 @@ -import torch.nn as nn -import numpy as np -from pytorch_lightning import LightningModule -from test_tube import HyperOptArgumentParser -from torchvision.datasets import MNIST -import torchvision.transforms as transforms -import torch -import torch.nn.functional as F - - -class ExampleModel1(LightningModule): - """ - Sample model to show how to define a template - """ - - def __init__(self, hparams): - # init superclass - super(ExampleModel1, self).__init__(hparams) - - self.batch_size = hparams.batch_size - - # build model - self.__build_model() - - # --------------------- - # MODEL SETUP - # --------------------- - def __build_model(self): - """ - Layout model - :return: - """ - self.c_d1 = nn.Linear(in_features=self.hparams.in_features, out_features=self.hparams.hidden_dim) - self.c_d1_bn = nn.BatchNorm1d(self.hparams.hidden_dim) - self.c_d1_drop = nn.Dropout(self.hparams.drop_prob) - - self.c_d2 = nn.Linear(in_features=self.hparams.hidden_dim, out_features=self.hparams.out_features) - - # --------------------- - # TRAINING - # --------------------- - def forward(self, x): - x = self.c_d1(x) - x = F.tanh(x) - x = self.c_d1_bn(x) - x = self.c_d1_drop(x) - - x = self.c_d2(x) - logits = F.log_softmax(x, dim=1) - - return logits - - def loss(self, labels, logits): - nll = F.nll_loss(logits, labels) - return nll - - def training_step(self, data_batch): - """ - Called inside the training loop - :param data_batch: - :return: - """ - # forward pass - x, y = data_batch - x = x.view(x.size(0), -1) - y_hat = self.forward(x) - - # calculate loss - loss_val = self.loss(y, y_hat) - - tqdm_dic = {'jefe': 1} - return loss_val, tqdm_dic - - def validation_step(self, data_batch): - """ - Called inside the validation loop - :param data_batch: - :return: - """ - x, y = data_batch - x = x.view(x.size(0), -1) - y_hat = self.forward(x) - - loss_val = self.loss(y, y_hat) - - # acc - labels_hat = torch.argmax(y_hat, dim=1) - val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0) - - output = {'y_hat': y_hat, 'val_loss': loss_val.item(), 'val_acc': val_acc} - return output - - def validation_end(self, outputs): - """ - Called at the end of validation to aggregate outputs - :param outputs: list of individual outputs of each validation step - :return: - """ - val_loss_mean = 0 - accs = [] - for output in outputs: - val_loss_mean += output['val_loss'] - accs.append(output['val_acc']) - - val_loss_mean /= len(outputs) - tqdm_dic = {'val_loss': val_loss_mean, 'val_acc': np.mean(accs)} - return tqdm_dic - - def update_tng_log_metrics(self, logs): - return logs - - # --------------------- - # MODEL SAVING - # --------------------- - def get_save_dict(self): - checkpoint = { - 'state_dict': self.state_dict(), - } - - return checkpoint - - def load_model_specific(self, checkpoint): - self.load_state_dict(checkpoint['state_dict']) - pass - - # --------------------- - # TRAINING SETUP - # --------------------- - def configure_optimizers(self): - """ - return whatever optimizers we want here - :return: list of optimizers - """ - optimizer = self.choose_optimizer(self.hparams.optimizer_name, self.parameters(), {'lr': self.hparams.learning_rate}, 'optimizer') - self.optimizers = [optimizer] - return self.optimizers - - def __dataloader(self, train): - # init data generators - transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))]) - - dataset = MNIST(root=self.hparams.data_root, train=train, transform=transform, download=True) - - loader = torch.utils.data.DataLoader( - dataset=dataset, - batch_size=self.hparams.batch_size, - shuffle=True - ) - - return loader - - @data_loader - def tng_dataloader(self): - if self._tng_dataloader is None: - try: - self._tng_dataloader = self.__dataloader(train=True) - except Exception as e: - print(e) - raise e - return self._tng_dataloader - - @property - def val_dataloader(self): - if self._val_dataloader is None: - try: - self._val_dataloader = self.__dataloader(train=False) - except Exception as e: - print(e) - raise e - return self._val_dataloader - - @property - def test_dataloader(self): - if self._test_dataloader is None: - try: - self._test_dataloader = self.__dataloader(train=False) - except Exception as e: - print(e) - raise e - return self._test_dataloader - - @staticmethod - def add_model_specific_args(parent_parser): - parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser]) - - # param overwrites - # parser.set_defaults(gradient_clip=5.0) - - # network params - parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False) - parser.add_argument('--in_features', default=28*28) - parser.add_argument('--hidden_dim', default=500) - parser.add_argument('--out_features', default=10) - - # data - parser.add_argument('--data_root', default='/Users/williamfalcon/Developer/personal/research_lib/research_proj/datasets/mnist', type=str) - - # training params (opt) - parser.opt_list('--learning_rate', default=0.001, type=float, options=[0.0001, 0.0005, 0.001, 0.005], - tunable=False) - parser.opt_list('--batch_size', default=256, type=int, options=[32, 64, 128, 256], tunable=False) - parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False) - return parser From 0ee034482049982d5b4fbc6be50cbca0c22c454e Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 21:39:53 -0400 Subject: [PATCH 05/24] removed old template --- .../lightning_module_template.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/pytorch_lightning/examples/new_project_templates/lightning_module_template.py b/pytorch_lightning/examples/new_project_templates/lightning_module_template.py index 0a4dab26..608e534e 100644 --- a/pytorch_lightning/examples/new_project_templates/lightning_module_template.py +++ b/pytorch_lightning/examples/new_project_templates/lightning_module_template.py @@ -154,17 +154,6 @@ class LightningTemplateModel(LightningModule): tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()} return tqdm_dic - # --------------------- - # MODEL SAVING - # --------------------- - def get_save_dict(self): - checkpoint = {'state_dict': self.state_dict()} - return checkpoint - - def load_model_specific(self, checkpoint): - self.load_state_dict(checkpoint['state_dict']) - pass - # --------------------- # TRAINING SETUP # --------------------- From 4148c36abddbec1d95c51b975cb9d4fc3cff238f Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 21:55:01 -0400 Subject: [PATCH 06/24] added model save load test --- pytorch_lightning/root_module/model_saving.py | 7 ++- tests/test_models.py | 56 +++++++++++++++++++ 2 files changed, 60 insertions(+), 3 deletions(-) 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 From 265411572fbab9a4b8afb92c9378c9effc22fc77 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:04:27 -0400 Subject: [PATCH 07/24] fixed hpc save, load. cleaned apu --- pytorch_lightning/root_module/model_saving.py | 37 ++++++++----------- 1 file changed, 16 insertions(+), 21 deletions(-) diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index e9224e2b..03cb5864 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -7,7 +7,7 @@ from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistr class ModelIO(object): - def load_model_specific(self, checkpoint): + def on_load_checkpoint(self, checkpoint): """ Do something with the checkpoint Gives model a chance to load something before state_dict is restored @@ -16,25 +16,24 @@ class ModelIO(object): """ pass - def get_save_dict(self): + def on_save_checkpoint(self, checkpoint): """ - Return specific things for the model - Called before trainer requests the state_dict - :return: + Give the model a chance to add something to the checkpoint. + state_dict is already there """ pass # ------------------------- # OPTIONAL HOOKS # ------------------------- - def on_hpc_save(self): + def on_hpc_save(self, checkpoint): """ Hook to do whatever you need right before Slurm manager saves the model :return: """ pass - def on_hpc_load(self): + def on_hpc_load(self, checkpoint): """ Hook to do whatever you need right before Slurm manager loads the model :return: @@ -78,13 +77,13 @@ class TrainerIO(object): checkpoint['optimizer_states'] = optimizer_states - # request what to save from the model - model = self.__get_model() - checkpoint_dict = model.get_save_dict() - checkpoint.update(checkpoint_dict) - # add the state_dict from the model - checkpoint['state_dict'] = checkpoint_dict + model = self.__get_model() + checkpoint['state_dict'] = model.get_state_dict + + # give the model a chance to add a few things + model.on_save_checkpoint(checkpoint) + return checkpoint # -------------------- @@ -153,13 +152,12 @@ class TrainerIO(object): # give model a chance to do something on hpc_save model = self.__get_model() - model.on_hpc_save() + checkpoint = self.dump_checkpoint() - # request what to save from the model - checkpoint_dict = self.dump_checkpoint() + model.on_hpc_save(checkpoint) # do the actual save - torch.save(checkpoint_dict, filepath) + torch.save(checkpoint, filepath) return filepath @@ -177,14 +175,11 @@ class TrainerIO(object): # load model state model = self.__get_model() - # give model a chance to load something - model.load_model_specific(checkpoint) - # load the state_dict on the model automatically model.load_state_dict(checkpoint['state_dict']) # call model hook - model.on_hpc_load() + model.on_hpc_load(checkpoint) def max_ckpt_in_folder(self, path): From 64de447545ba7d17da61c321c4d182bb2d2339dc Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:07:02 -0400 Subject: [PATCH 08/24] fixed hpc save, load. cleaned apu --- pytorch_lightning/root_module/root_module.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py index a25f7f6f..b49dd3af 100644 --- a/pytorch_lightning/root_module/root_module.py +++ b/pytorch_lightning/root_module/root_module.py @@ -108,13 +108,12 @@ class LightningModule(GradInformation, ModelIO, ModelHooks): else: checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage) + # load the state_dict on the model automatically model = cls(hparams) + model.load_state_dict(checkpoint['state_dict']) # give model a chance to load something - model.load_model_specific(checkpoint) - - # load the state_dict on the model automatically - model.load_state_dict(checkpoint['state_dict']) + model.on_load_checkpoint(checkpoint) return model From 348223a702bfe4724e9ec8ac46050e46b445eba5 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:09:35 -0400 Subject: [PATCH 09/24] fixed hpc save, load. cleaned apu --- pytorch_lightning/root_module/model_saving.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index 03cb5864..142d2b33 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -79,7 +79,7 @@ class TrainerIO(object): # add the state_dict from the model model = self.__get_model() - checkpoint['state_dict'] = model.get_state_dict + checkpoint['state_dict'] = model.state_dict() # give the model a chance to add a few things model.on_save_checkpoint(checkpoint) From a6ae97ac0922382a1e2ed7c593bb064f141001a3 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:13:06 -0400 Subject: [PATCH 10/24] fixed hpc save, load. cleaned apu --- tests/test_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index 1a315625..bdd31a5c 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -115,7 +115,7 @@ def test_model_saving_loading(): # make prediction # assert that both predictions are the same new_pred = model_2(x) - assert torch.eq(pred_before_saving, new_pred) + assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1 clear_save_dir() From c61e13f0ffa13d11b66a98538bf74ddbf655b179 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:13:41 -0400 Subject: [PATCH 11/24] fixed hpc save, load. cleaned apu --- tests/test_models.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_models.py b/tests/test_models.py index bdd31a5c..5ead6de8 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -66,6 +66,10 @@ def test_dp_output_reduce(): def test_model_saving_loading(): + """ + Tests use case where trainer saves the model, and user loads it from tags independently + :return: + """ hparams = get_hparams() model = LightningTestModel(hparams) From b5419fcd8b11580744376ef9d8ed91d510298508 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:24:01 -0400 Subject: [PATCH 12/24] added clean slurm save load test --- tests/test_models.py | 44 +++++++++++++++++++++++++++----------------- 1 file changed, 27 insertions(+), 17 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 5ead6de8..768acd05 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -90,7 +90,6 @@ def test_model_saving_loading(): # 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' @@ -124,7 +123,7 @@ def test_model_saving_loading(): clear_save_dir() -def test_cpu_slurm_saving_loading(): +def test_cpu_slurm_save_load(): """ Verify model save/load/checkpoint on CPU :return: @@ -154,38 +153,49 @@ def test_cpu_slurm_saving_loading(): # traning complete assert result == 1, 'amp + ddp model failed to complete' - # test saving checkpoint - ckpt_test = os.path.join(save_dir, 'test.ckpt') - trainer.save_checkpoint(ckpt_test) + # predict with trained model before saving + # make a prediction + for batch in model.test_dataloader: + break + + x, y = batch + x = x.view(x.size(0), -1) + + model.eval() + pred_before_saving = model(x) # test registering a save function trainer.enable_auto_hpc_walltime_manager() - # test model loading with a map_location - pretrained_model = load_model(exp, save_dir, True) - - # test model preds - run_prediction(model.test_dataloader, pretrained_model) - - trainer.model = pretrained_model - trainer.optimizers = pretrained_model.configure_optimizers() - # test HPC saving + # simulate snapshot on slurm saved_filepath = trainer.hpc_save(save_dir, exp) assert os.path.exists(saved_filepath) + # wipe-out trainer model + # we want to see if the weights come back correctly + trainer.model = LightningTestModel(hparams) + # test HPC loading trainer.global_step = 20000000 trainer.hpc_load(save_dir, on_gpu=False) assert trainer.global_step == real_global_step and trainer.global_step != 20000000 - # test freeze on gpu - model.freeze() - model.unfreeze() + # predict with loaded model to make sure answers are the same + new_pred = trainer.model(x) + assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1 clear_save_dir() +def test_model_freeze_unfreeze(): + hparams = get_hparams() + model = LightningTestModel(hparams) + + model.freeze() + model.unfreeze() + + def test_amp_gpu_ddp_slurm_managed(): """ Make sure DDP + AMP work From ffa7a0dbab42df2e0dcadde1ee5b4a519ae80481 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:26:55 -0400 Subject: [PATCH 13/24] added clean slurm save load test --- pytorch_lightning/root_module/model_saving.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index 142d2b33..361a0854 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -181,6 +181,7 @@ class TrainerIO(object): # call model hook model.on_hpc_load(checkpoint) + self.model = model def max_ckpt_in_folder(self, path): files = os.listdir(path) From 57edb08bd8a73c01e326e4bc259f2ece41acc892 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:28:09 -0400 Subject: [PATCH 14/24] added clean slurm save load test --- pytorch_lightning/root_module/model_saving.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index 361a0854..438d67ea 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -176,12 +176,12 @@ class TrainerIO(object): model = self.__get_model() # load the state_dict on the model automatically + pdb.set_trace() model.load_state_dict(checkpoint['state_dict']) # call model hook model.on_hpc_load(checkpoint) - self.model = model def max_ckpt_in_folder(self, path): files = os.listdir(path) From f1de62671de7cff5f234b791b56210ef72900cfd Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:32:27 -0400 Subject: [PATCH 15/24] added clean slurm save load test --- pytorch_lightning/root_module/model_saving.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index 438d67ea..0638f294 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -182,7 +182,6 @@ class TrainerIO(object): # call model hook model.on_hpc_load(checkpoint) - def max_ckpt_in_folder(self, path): files = os.listdir(path) files = [x for x in files if 'ckpt_' in x] From f5a01edfb8112e4c10e4aa6d2a8683e6a0837ea7 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:32:34 -0400 Subject: [PATCH 16/24] added clean slurm save load test --- pytorch_lightning/root_module/model_saving.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index 0638f294..c5831317 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -176,7 +176,6 @@ class TrainerIO(object): model = self.__get_model() # load the state_dict on the model automatically - pdb.set_trace() model.load_state_dict(checkpoint['state_dict']) # call model hook From 8e3a0443c73c9b28a685c69ed59e28a7934023c9 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:33:00 -0400 Subject: [PATCH 17/24] added clean slurm save load test --- tests/test_models.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_models.py b/tests/test_models.py index 768acd05..0bf8d350 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -182,6 +182,7 @@ def test_cpu_slurm_save_load(): assert trainer.global_step == real_global_step and trainer.global_step != 20000000 # predict with loaded model to make sure answers are the same + trainer.model.eval() new_pred = trainer.model(x) assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1 From 2a4081e5370cb9748847994bdf7360efe2ef3579 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:33:31 -0400 Subject: [PATCH 18/24] added clean slurm save load test --- tests/test_models.py | 131 ++++++++++++++++++++++--------------------- 1 file changed, 67 insertions(+), 64 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 0bf8d350..c83ff428 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -24,6 +24,73 @@ np.random.seed(SEED) # ------------------------------------------------------------------------ # TESTS # ------------------------------------------------------------------------ + +def test_cpu_slurm_save_load(): + """ + Verify model save/load/checkpoint on CPU + :return: + """ + 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' + + # predict with trained model before saving + # make a prediction + for batch in model.test_dataloader: + break + + x, y = batch + x = x.view(x.size(0), -1) + + model.eval() + pred_before_saving = model(x) + + # test registering a save function + trainer.enable_auto_hpc_walltime_manager() + + # test HPC saving + # simulate snapshot on slurm + saved_filepath = trainer.hpc_save(save_dir, exp) + assert os.path.exists(saved_filepath) + + # wipe-out trainer model + # we want to see if the weights come back correctly + trainer.model = LightningTestModel(hparams) + + # test HPC loading + trainer.global_step = 20000000 + trainer.hpc_load(save_dir, on_gpu=False) + assert trainer.global_step == real_global_step and trainer.global_step != 20000000 + + # predict with loaded model to make sure answers are the same + trainer.model.eval() + new_pred = trainer.model(x) + assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1 + + clear_save_dir() + + def test_loading_meta_tags(): hparams = get_hparams() @@ -123,70 +190,6 @@ def test_model_saving_loading(): clear_save_dir() -def test_cpu_slurm_save_load(): - """ - Verify model save/load/checkpoint on CPU - :return: - """ - 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' - - # predict with trained model before saving - # make a prediction - for batch in model.test_dataloader: - break - - x, y = batch - x = x.view(x.size(0), -1) - - model.eval() - pred_before_saving = model(x) - - # test registering a save function - trainer.enable_auto_hpc_walltime_manager() - - # test HPC saving - # simulate snapshot on slurm - saved_filepath = trainer.hpc_save(save_dir, exp) - assert os.path.exists(saved_filepath) - - # wipe-out trainer model - # we want to see if the weights come back correctly - trainer.model = LightningTestModel(hparams) - - # test HPC loading - trainer.global_step = 20000000 - trainer.hpc_load(save_dir, on_gpu=False) - assert trainer.global_step == real_global_step and trainer.global_step != 20000000 - - # predict with loaded model to make sure answers are the same - trainer.model.eval() - new_pred = trainer.model(x) - assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1 - - clear_save_dir() def test_model_freeze_unfreeze(): From 322436519096b65585424098683517bd2fdc3035 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:39:44 -0400 Subject: [PATCH 19/24] added clean slurm save load test --- tests/test_models.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index c83ff428..57fe9f77 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -74,8 +74,9 @@ def test_cpu_slurm_save_load(): saved_filepath = trainer.hpc_save(save_dir, exp) assert os.path.exists(saved_filepath) - # wipe-out trainer model + # wipe-out trainer and model # we want to see if the weights come back correctly + trainer = Trainer(**trainer_options) trainer.model = LightningTestModel(hparams) # test HPC loading From 61c82611eb33786c9d0b8f7a55ae294ea143410d Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:40:07 -0400 Subject: [PATCH 20/24] added clean slurm save load test --- tests/test_models.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 57fe9f77..2977f460 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -80,9 +80,8 @@ def test_cpu_slurm_save_load(): trainer.model = LightningTestModel(hparams) # test HPC loading - trainer.global_step = 20000000 trainer.hpc_load(save_dir, on_gpu=False) - assert trainer.global_step == real_global_step and trainer.global_step != 20000000 + assert trainer.global_step == real_global_step and trainer.global_step > 0 # predict with loaded model to make sure answers are the same trainer.model.eval() From f183ac2a1c71dce6c685058a97d26e51b580c112 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:51:33 -0400 Subject: [PATCH 21/24] added clean slurm save load test --- pytorch_lightning/models/trainer.py | 1 - tests/test_models.py | 34 ++++++++++++++++++++--------- 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 44cafb0c..9c4b7d5f 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -161,7 +161,6 @@ class Trainer(TrainerIO): self.nb_tng_batches = None self.nb_test_batches = None - # gpus come in as a string. # if gpus = -1 then use all available devices # otherwise, split the string using commas diff --git a/tests/test_models.py b/tests/test_models.py index 2977f460..5cee2e5e 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -75,9 +75,16 @@ def test_cpu_slurm_save_load(): assert os.path.exists(saved_filepath) # wipe-out trainer and model + # retrain with not much data... this simulates picking training back up after slurm # we want to see if the weights come back correctly + continue_tng_hparams = get_hparams(continue_training=True) + trainer_options = dict( + max_nb_epochs=1, + cluster=SlurmCluster(continue_tng_hparams), + experiment=exp, + checkpoint_callback=ModelCheckpoint(save_dir) + ) trainer = Trainer(**trainer_options) - trainer.model = LightningTestModel(hparams) # test HPC loading trainer.hpc_load(save_dir, on_gpu=False) @@ -568,16 +575,23 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True): clear_save_dir() -def get_hparams(): +def get_hparams(continue_training=False): root_dir = os.path.dirname(os.path.realpath(__file__)) - hparams = Namespace(**{'drop_prob': 0.2, - 'batch_size': 32, - 'in_features': 28*28, - 'learning_rate': 0.001*8, - 'optimizer_name': 'adam', - 'data_root': os.path.join(root_dir, 'mnist'), - 'out_features': 10, - 'hidden_dim': 1000}) + + args = { + 'drop_prob': 0.2, + 'batch_size': 32, + 'in_features': 28*28, + 'learning_rate': 0.001*8, + 'optimizer_name': 'adam', + 'data_root': os.path.join(root_dir, 'mnist'), + 'out_features': 10, + 'hidden_dim': 1000} + + if continue_training: + args['test_tube_do_checkpoint_load'] = True + + hparams = Namespace(**args) return hparams From 53b781709e058e27da1543069999ce107fc3f89b Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 22:57:49 -0400 Subject: [PATCH 22/24] added clean slurm save load test --- pytorch_lightning/models/trainer.py | 6 +++++- tests/test_models.py | 21 ++++++++++++++------- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 9c4b7d5f..ff11150b 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -610,14 +610,18 @@ class Trainer(TrainerIO): if self.proc_rank == 0: self.experiment.save() + # track model now. + # if cluster resets state, the model will update with the saved weights + self.model = model + # enable cluster checkpointing + # also restores training state if self.cluster is not None: # pragma: no cover self.enable_auto_hpc_walltime_manager() # --------------------------- # CORE TRAINING LOOP # --------------------------- - self.model = model self.__train() def __train(self): diff --git a/tests/test_models.py b/tests/test_models.py index 5cee2e5e..e95870e8 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -85,15 +85,22 @@ def test_cpu_slurm_save_load(): checkpoint_callback=ModelCheckpoint(save_dir) ) trainer = Trainer(**trainer_options) + model = LightningTestModel(hparams) - # test HPC loading - trainer.hpc_load(save_dir, on_gpu=False) - assert trainer.global_step == real_global_step and trainer.global_step > 0 + # set the epoch start hook so we can predict before the model does the full training + def assert_pred_same(): + assert trainer.global_step == real_global_step and trainer.global_step > 0 - # predict with loaded model to make sure answers are the same - trainer.model.eval() - new_pred = trainer.model(x) - assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1 + # predict with loaded model to make sure answers are the same + trainer.model.eval() + new_pred = trainer.model(x) + assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1 + + model.on_epoch_start = assert_pred_same + + # by calling fit again, we trigger training, loading weights from the cluster + # and our hook to predict using current model before any more weight updates + trainer.fit(model) clear_save_dir() From 64586f271d6ee9c1a771ca21297a6c39b77095c4 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 23:02:18 -0400 Subject: [PATCH 23/24] added clean slurm save load test --- tests/test_models.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index e95870e8..6540204a 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -40,9 +40,10 @@ def test_cpu_slurm_save_load(): exp.argparse(hparams) exp.save() + cluster_a = SlurmCluster() trainer_options = dict( max_nb_epochs=1, - cluster=SlurmCluster(), + cluster=cluster_a, experiment=exp, checkpoint_callback=ModelCheckpoint(save_dir) ) @@ -82,7 +83,8 @@ def test_cpu_slurm_save_load(): max_nb_epochs=1, cluster=SlurmCluster(continue_tng_hparams), experiment=exp, - checkpoint_callback=ModelCheckpoint(save_dir) + checkpoint_callback=ModelCheckpoint(save_dir), + hpc_exp_number=cluster_a.hpc_exp_number ) trainer = Trainer(**trainer_options) model = LightningTestModel(hparams) From 587c195298171c8e27e404da8998fda4415b53b8 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 26 Jul 2019 23:04:41 -0400 Subject: [PATCH 24/24] added clean slurm save load test --- tests/test_models.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 6540204a..e8fa339b 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -78,13 +78,12 @@ def test_cpu_slurm_save_load(): # wipe-out trainer and model # retrain with not much data... this simulates picking training back up after slurm # we want to see if the weights come back correctly - continue_tng_hparams = get_hparams(continue_training=True) + continue_tng_hparams = get_hparams(continue_training=True, hpc_exp_number=cluster_a.hpc_exp_number) trainer_options = dict( max_nb_epochs=1, cluster=SlurmCluster(continue_tng_hparams), experiment=exp, checkpoint_callback=ModelCheckpoint(save_dir), - hpc_exp_number=cluster_a.hpc_exp_number ) trainer = Trainer(**trainer_options) model = LightningTestModel(hparams) @@ -584,7 +583,7 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True): clear_save_dir() -def get_hparams(continue_training=False): +def get_hparams(continue_training=False, hpc_exp_number=0): root_dir = os.path.dirname(os.path.realpath(__file__)) args = { @@ -599,6 +598,7 @@ def get_hparams(continue_training=False): if continue_training: args['test_tube_do_checkpoint_load'] = True + args['hpc_exp_number'] = hpc_exp_number hparams = Namespace(**args) return hparams