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 # --------------------- 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 diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 44cafb0c..ff11150b 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 @@ -611,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/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index 25377851..c5831317 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -4,34 +4,36 @@ import re import pdb from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel + 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 :param checkpoint: :return: """ - raise NotImplementedError + pass - def get_save_dict(self): + def on_save_checkpoint(self, checkpoint): """ - Return specific things for the model - :return: + Give the model a chance to add something to the checkpoint. + state_dict is already there """ - raise NotImplementedError + 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: @@ -75,12 +77,13 @@ class TrainerIO(object): checkpoint['optimizer_states'] = optimizer_states - # request what to save from the model + # add the state_dict from the model model = self.__get_model() - checkpoint_dict = model.get_save_dict() + checkpoint['state_dict'] = model.state_dict() + + # give the model a chance to add a few things + model.on_save_checkpoint(checkpoint) - # merge trainer and model saving items - checkpoint.update(checkpoint_dict) return checkpoint # -------------------- @@ -149,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 @@ -167,15 +169,17 @@ 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() - 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): files = os.listdir(path) diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py index d6d740f0..b49dd3af 100644 --- a/pytorch_lightning/root_module/root_module.py +++ b/pytorch_lightning/root_module/root_module.py @@ -108,11 +108,13 @@ 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.on_load_checkpoint(checkpoint) - # allow model to load - model.load_model_specific(checkpoint) - model.load_state_dict(checkpoint['state_dict'], strict=False) return model def summarize(self): 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 # --------------------- diff --git a/tests/test_models.py b/tests/test_models.py index e5781d32..e8fa339b 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -24,6 +24,88 @@ 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() + + cluster_a = SlurmCluster() + trainer_options = dict( + max_nb_epochs=1, + cluster=cluster_a, + 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 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, 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), + ) + trainer = Trainer(**trainer_options) + model = LightningTestModel(hparams) + + # 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 + + 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() + + def test_loading_meta_tags(): hparams = get_hparams() @@ -43,6 +125,7 @@ def test_loading_meta_tags(): clear_save_dir() + def test_dp_output_reduce(): # test identity when we have a single gpu @@ -64,9 +147,9 @@ def test_dp_output_reduce(): assert reduced['b']['c'] == out['b']['c'] -def test_cpu_slurm_saving_loading(): +def test_model_saving_loading(): """ - Verify model save/load/checkpoint on CPU + Tests use case where trainer saves the model, and user loads it from tags independently :return: """ hparams = get_hparams() @@ -89,43 +172,49 @@ def test_cpu_slurm_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' - # test saving checkpoint - ckpt_test = os.path.join(save_dir, 'test.ckpt') - trainer.save_checkpoint(ckpt_test) + # make a prediction + for batch in model.test_dataloader: + break - # test registering a save function - trainer.enable_auto_hpc_walltime_manager() + x, y = batch + x = x.view(x.size(0), -1) - # test model loading with a map_location - pretrained_model = load_model(exp, save_dir, True) + # generate preds before saving model + model.eval() + pred_before_saving = model(x) - # test model preds - run_prediction(model.test_dataloader, pretrained_model) + # save model + new_weights_path = os.path.join(save_dir, 'save_test.ckpt') + trainer.save_checkpoint(new_weights_path) - trainer.model = pretrained_model - trainer.optimizers = pretrained_model.configure_optimizers() + # 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() - # test HPC saving - saved_filepath = trainer.hpc_save(save_dir, exp) - assert os.path.exists(saved_filepath) - - # 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() + # make prediction + # assert that both predictions are the same + new_pred = model_2(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 @@ -494,16 +583,24 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True): clear_save_dir() -def get_hparams(): +def get_hparams(continue_training=False, hpc_exp_number=0): 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 + args['hpc_exp_number'] = hpc_exp_number + + hparams = Namespace(**args) return hparams