From 2a04be038624bd8867b01b4b7bfe967b010fb941 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Mon, 2 Mar 2020 17:12:22 -0500 Subject: [PATCH] No auto load weights (#985) * remove autoload * remove autoload * added weights loading docs * checkpoint loading saving docs * checkpoint loading saving docs * checkpoint loading saving docs * docs (#1010) * remove autoload * remove autoload * added weights loading docs * checkpoint loading saving docs * checkpoint loading saving docs * checkpoint loading saving docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs --- docs/source/checkpointing.rst | 80 ---------- docs/source/index.rst | 2 +- docs/source/weights_loading.rst | 76 ++++++++++ pytorch_lightning/trainer/callback_config.py | 11 +- pytorch_lightning/trainer/evaluation_loop.py | 4 + pytorch_lightning/trainer/trainer.py | 9 +- pytorch_lightning/trainer/training_io.py | 150 ++----------------- pytorch_lightning/trainer/training_loop.py | 31 ---- tests/models/utils.py | 3 +- tests/test_restore_models.py | 71 +-------- 10 files changed, 112 insertions(+), 325 deletions(-) delete mode 100644 docs/source/checkpointing.rst create mode 100644 docs/source/weights_loading.rst diff --git a/docs/source/checkpointing.rst b/docs/source/checkpointing.rst deleted file mode 100644 index 6ec85e8a..00000000 --- a/docs/source/checkpointing.rst +++ /dev/null @@ -1,80 +0,0 @@ -Checkpointing -============== - -.. _model-saving: - -Model saving -------------------- -To save a LightningModule, provide a :meth:`pytorch_lightning.callbacks.ModelCheckpoint` callback. - -The Lightning checkpoint also saves the hparams (hyperparams) passed into the LightningModule init. - -.. note:: hparams is a `Namespace `_ or dictionary. - -.. code-block:: python - :emphasize-lines: 8 - - from argparse import Namespace - - # usually these come from command line args - args = Namespace(**{'learning_rate':0.001}) - - # define you module to have hparams as the first arg - # this means your checkpoint will have everything that went into making - # this model (in this case, learning rate) - class MyLightningModule(pl.LightningModule): - - def __init__(self, hparams, ...): - self.hparams = hparams - - my_model = MyLightningModule(args) - - # auto-saves checkpoint - checkpoint_callback = ModelCheckpoint(filepath='my_path') - Trainer(checkpoint_callback=checkpoint_callback) - - -Model loading ------------------------------------ - -To load a model, use :meth:`pytorch_lightning.core.LightningModule.load_from_checkpoint` - -.. note:: If lightning created your checkpoint, your model will receive all the hyperparameters used - to create the checkpoint. (See: :ref:`model-saving`). - -.. code-block:: python - - # load weights without mapping - MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt') - - # load weights mapping all weights from GPU 1 to GPU 0 - map_location = {'cuda:1':'cuda:0'} - MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt', map_location=map_location) - -Restoring training session ------------------------------------ - -If you want to pick up training from where you left off, you have a few options. - -1. Pass in a logger with the same experiment version to continue training. - -.. code-block:: python - - # train the first time and set the version number - logger = TensorboardLogger(version=10) - trainer = Trainer(logger=logger) - trainer.fit(model) - - # when you init another logger with that same version, the model - # will continue where it left off - logger = TensorboardLogger(version=10) - trainer = Trainer(logger=logger) - trainer.fit(model) - -2. A second option is to pass in a path to a checkpoint (see: :ref:`pytorch_lightning.trainer.trainer.Trainer`). - -.. code-block:: python - - # train the first time and set the version number - trainer = Trainer(resume_from_checkpoint='some/path/to/my_checkpoint.ckpt') - trainer.fit(model) \ No newline at end of file diff --git a/docs/source/index.rst b/docs/source/index.rst index 9eb6cc7b..d44c80ef 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -45,7 +45,6 @@ PyTorch-Lightning Documentation :caption: Common Use Cases apex - checkpointing slurm debugging experiment_logging @@ -54,6 +53,7 @@ PyTorch-Lightning Documentation fast_training hooks multi_gpu + weights_loading single_gpu sequences training_tricks diff --git a/docs/source/weights_loading.rst b/docs/source/weights_loading.rst new file mode 100644 index 00000000..21a24ef7 --- /dev/null +++ b/docs/source/weights_loading.rst @@ -0,0 +1,76 @@ +Saving and loading weights +========================== + +Lightning can automate saving and loading checkpoints. + +Checkpoint saving +----------------- + +Checkpointing is enabled by default to the current working directory. +To change the checkpoint path pass in: + +.. code-block:: python + + Trainer(default_save_path='/your/path/to/save/checkpoints') + +To modify the behavior of checkpointing pass in your own callback. + +.. code-block:: python + + from pytorch_lightning.callbacks import ModelCheckpoint + + # DEFAULTS used by the Trainer + checkpoint_callback = ModelCheckpoint( + filepath=os.getcwd(), + save_best_only=True, + verbose=True, + monitor='val_loss', + mode='min', + prefix='' + ) + + trainer = Trainer(checkpoint_callback=checkpoint_callback) + + +Or disable it by passing + +.. code-block:: python + + trainer = Trainer(checkpoint_callback=False) + + +The Lightning checkpoint also saves the hparams (hyperparams) passed into the LightningModule init. + +.. note:: hparams is a `Namespace `_. + +.. code-block:: python + :emphasize-lines: 8 + + from argparse import Namespace + + # usually these come from command line args + args = Namespace(**{'learning_rate':0.001}) + + # define you module to have hparams as the first arg + # this means your checkpoint will have everything that went into making + # this model (in this case, learning rate) + class MyLightningModule(pl.LightningModule): + + def __init__(self, hparams, ...): + self.hparams = hparams + +Checkpoint Loading +------------------ + +You might want to not only load a model but also continue training it. Use this method to +restore the trainer state as well. This will continue from the epoch and global step you last left off. +However, the dataloaders will start from the first batch again (if you shuffled it shouldn't matter). + +.. code-block:: python + + model = MyLightingModule.load_from_checkpoint(PATH) + model.eval() + y_hat = model(x) + +A LightningModule is no different than a nn.Module. This means you can load it and use it for +predictions as you would a nn.Module. \ No newline at end of file diff --git a/pytorch_lightning/trainer/callback_config.py b/pytorch_lightning/trainer/callback_config.py index 8a17698e..140c8f74 100644 --- a/pytorch_lightning/trainer/callback_config.py +++ b/pytorch_lightning/trainer/callback_config.py @@ -62,7 +62,7 @@ class TrainerCallbackConfigMixin(ABC): self.weights_save_path = self.default_save_path def configure_early_stopping(self, early_stop_callback): - if early_stop_callback is True: + if early_stop_callback is True or None: self.early_stop_callback = EarlyStopping( monitor='val_loss', patience=3, @@ -71,15 +71,6 @@ class TrainerCallbackConfigMixin(ABC): mode='min' ) self.enable_early_stop = True - elif early_stop_callback is None: - self.early_stop_callback = EarlyStopping( - monitor='val_loss', - patience=3, - strict=False, - verbose=False, - mode='min' - ) - self.enable_early_stop = True elif not early_stop_callback: self.early_stop_callback = None self.enable_early_stop = False diff --git a/pytorch_lightning/trainer/evaluation_loop.py b/pytorch_lightning/trainer/evaluation_loop.py index a00a228b..b713a744 100644 --- a/pytorch_lightning/trainer/evaluation_loop.py +++ b/pytorch_lightning/trainer/evaluation_loop.py @@ -347,6 +347,10 @@ class TrainerEvaluationLoopMixin(ABC): # add metrics to prog bar self.add_tqdm_metrics(prog_bar_metrics) + # log results of test + if test_mode: + model.print(prog_bar_metrics) + # log metrics self.log_metrics(log_metrics, {}) diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index 6d5811c2..9a5920a3 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -1096,7 +1096,8 @@ class Trainer(TrainerIOMixin, self.register_slurm_signal_handlers() # print model summary - if self.proc_rank == 0 and self.weights_summary is not None: + # TODO: remove self.testing condition because model.summarize() is wiping out the weights + if self.proc_rank == 0 and self.weights_summary is not None and not self.testing: if self.weights_summary in ['full', 'top']: ref_model.summarize(mode=self.weights_summary) else: @@ -1116,7 +1117,7 @@ class Trainer(TrainerIOMixin, # when testing requested only run test and return if self.testing: # only load test dataloader for testing - self.reset_test_dataloader(ref_model) + # self.reset_test_dataloader(ref_model) self.run_evaluation(test_mode=True) return @@ -1189,8 +1190,10 @@ class Trainer(TrainerIOMixin, """ self.testing = True if model is not None: + self.model = model self.fit(model) - self.run_evaluation(test_mode=True) + else: + self.run_evaluation(test_mode=True) class _PatchDataLoader(object): diff --git a/pytorch_lightning/trainer/training_io.py b/pytorch_lightning/trainer/training_io.py index 9cfe0537..b4592923 100644 --- a/pytorch_lightning/trainer/training_io.py +++ b/pytorch_lightning/trainer/training_io.py @@ -1,94 +1,3 @@ -""" -Lightning can automate saving and loading checkpoints -===================================================== - -Checkpointing is enabled by default to the current working directory. -To change the checkpoint path pass in:: - - Trainer(default_save_path='/your/path/to/save/checkpoints') - - -To modify the behavior of checkpointing pass in your own callback. - -.. code-block:: python - - from pytorch_lightning.callbacks import ModelCheckpoint - - # DEFAULTS used by the Trainer - checkpoint_callback = ModelCheckpoint( - filepath=os.getcwd(), - save_best_only=True, - verbose=True, - monitor='val_loss', - mode='min', - prefix='' - ) - - trainer = Trainer(checkpoint_callback=checkpoint_callback) - - -Restoring training session --------------------------- - -You might want to not only load a model but also continue training it. Use this method to -restore the trainer state as well. This will continue from the epoch and global step you last left off. -However, the dataloaders will start from the first batch again (if you shuffled it shouldn't matter). - -Lightning will restore the session if you pass a logger with the same version and there's a saved checkpoint. - -.. code-block:: python - - from pytorch_lightning import Trainer - from pytorch_lightning.loggers import TestTubeLogger - - logger = TestTubeLogger( - save_dir='./savepath', - version=1 # An existing version with a saved checkpoint - ) - trainer = Trainer( - logger=logger, - default_save_path='./savepath' - ) - - # this fit call loads model weights and trainer state - # the trainer continues seamlessly from where you left off - # without having to do anything else. - trainer.fit(model) - - -The trainer restores: - -- global_step -- current_epoch -- All optimizers -- All lr_schedulers -- Model weights - -You can even change the logic of your model as long as the weights and "architecture" of -the system isn't different. If you add a layer, for instance, it might not work. - -At a rough level, here's what happens inside Trainer :py:mod:`pytorch_lightning.base_module.model_saving.py`: - -.. code-block:: python - - self.global_step = checkpoint['global_step'] - self.current_epoch = checkpoint['epoch'] - - # restore the optimizers - optimizer_states = checkpoint['optimizer_states'] - for optimizer, opt_state in zip(self.optimizers, optimizer_states): - optimizer.load_state_dict(opt_state) - - # restore the lr schedulers - lr_schedulers = checkpoint['lr_schedulers'] - for scheduler, lrs_state in zip(self.lr_schedulers, lr_schedulers): - scheduler.load_state_dict(lrs_state) - - # uses the model you passed into trainer - model.load_state_dict(checkpoint['state_dict']) - -""" - import logging as log import os import re @@ -150,11 +59,11 @@ class TrainerIOMixin(ABC): # -------------------- def restore_weights(self, model): """ - To restore weights we have two cases. - First, attempt to restore hpc weights. If successful, don't restore - other weights. + We attempt to restore weights in this order: + 1. HPC weights. + 2. if no HPC weights restore checkpoint_path weights + 3. otherwise don't restore weights - Otherwise, try to restore actual weights :param model: :return: """ @@ -172,9 +81,6 @@ class TrainerIOMixin(ABC): if not did_restore_hpc_weights: if self.resume_from_checkpoint is not None: self.restore(self.resume_from_checkpoint, on_gpu=self.on_gpu) - else: - # restore weights if same exp version - self.restore_state_if_checkpoint_exists(model) # wait for all models to restore weights if self.use_ddp or self.use_ddp2: @@ -190,42 +96,6 @@ class TrainerIOMixin(ABC): if self.on_gpu: torch.cuda.empty_cache() - def restore_state_if_checkpoint_exists(self, model): - did_restore = False - - # do nothing if there's not dir or callback - no_ckpt_callback = (self.checkpoint_callback is None) or (not self.checkpoint_callback) - if no_ckpt_callback or not os.path.exists(self.checkpoint_callback.filepath): - return did_restore - - # restore trainer state and model if there is a weight for this experiment - last_epoch = -1 - last_ckpt_name = None - - # find last epoch - checkpoints = os.listdir(self.checkpoint_callback.filepath) - for name in checkpoints: - # ignore hpc ckpts - if 'hpc_' in name: - continue - - if '.ckpt' in name: - epoch = name.split('epoch_')[1] - epoch = int(re.sub('[^0-9]', '', epoch)) - - if epoch > last_epoch: - last_epoch = epoch - last_ckpt_name = name - - # restore last checkpoint - if last_ckpt_name is not None: - last_ckpt_path = os.path.join(self.checkpoint_callback.filepath, last_ckpt_name) - self.restore(last_ckpt_path, self.on_gpu) - log.info(f'Model and Trainer restored from checkpoint: {last_ckpt_path}') - did_restore = True - - return did_restore - # -------------------- # HPC SIGNAL HANDLING # -------------------- @@ -304,6 +174,18 @@ class TrainerIOMixin(ABC): self._atomic_save(checkpoint, filepath) def restore(self, checkpoint_path, on_gpu): + """ + Restore training state from checkpoint. + Also restores all training state like: + - epoch + - callbacks + - schedulers + - optimizer + :param checkpoint_path: + :param on_gpu: + + :return: + """ # if on_gpu: # checkpoint = torch.load(checkpoint_path) diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py index 03df1851..e166f6e9 100644 --- a/pytorch_lightning/trainer/training_loop.py +++ b/pytorch_lightning/trainer/training_loop.py @@ -25,37 +25,6 @@ It can be useful to force training for a minimum number of epochs or limit to a # DEFAULT trainer = Trainer(min_epochs=1, max_epochs=1000) -Early stopping --------------- - -The trainer already sets up default early stopping for you. -To modify this behavior, pass in your own EarlyStopping callback. - -.. code-block:: python - - from pytorch_lightning.callbacks import EarlyStopping - - # DEFAULTS used by Trainer - early_stop_callback = EarlyStopping( - monitor='val_loss', - min_delta=0.00, - patience=3, - verbose=False, - mode='min' - ) - - # without passing anything in, uses the default callback above - trainer = Trainer() - - # pass in your own to override the default callback - trainer = Trainer(early_stop_callback=early_stop_callback) - - # pass in min_epochs to enable the callback after min_epochs have run - trainer = Trainer(early_stop_callback=early_stop_callback, min_epochs=5) - - # pass in None to disable it - trainer = Trainer(early_stop_callback=None) - Force disable early stop ------------------------ diff --git a/tests/models/utils.py b/tests/models/utils.py index 7eceb100..df3c0411 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -158,8 +158,7 @@ def load_model(exp, root_weights_dir, module_class=LightningTemplateModel, path_ checkpoints = [x for x in os.listdir(root_weights_dir) if '.ckpt' in x] weights_dir = os.path.join(root_weights_dir, checkpoints[0]) - trained_model = module_class.load_from_metrics(weights_path=weights_dir, - tags_csv=tags_path) + trained_model = module_class.load_from_checkpoint(weights_dir) assert trained_model is not None, 'loading model failed' diff --git a/tests/test_restore_models.py b/tests/test_restore_models.py index 1ed36c2a..19acd1fd 100644 --- a/tests/test_restore_models.py +++ b/tests/test_restore_models.py @@ -126,6 +126,7 @@ def test_load_model_from_checkpoint(tmpdir): # fit model trainer = Trainer(**trainer_options) result = trainer.fit(model) + trainer.test() # correct result and ok accuracy assert result == 1, 'training failed to complete' @@ -140,6 +141,10 @@ def test_load_model_from_checkpoint(tmpdir): for k, v in vars(hparams).items(): assert getattr(pretrained_model.hparams, k) == v + # assert weights are the same + for (old_name, old_p), (new_name, new_p) in zip(model.named_parameters(), pretrained_model.named_parameters()): + assert torch.all(torch.eq(old_p, new_p)), 'loaded weights are not the same as the saved weights' + new_trainer = Trainer(**trainer_options) new_trainer.test(pretrained_model) @@ -203,7 +208,7 @@ def test_dp_resume(tmpdir): trainer_options = dict( show_progress_bar=True, - max_epochs=2, + max_epochs=3, gpus=2, distributed_backend='dp', ) @@ -240,7 +245,7 @@ def test_dp_resume(tmpdir): new_logger = tutils.get_test_tube_logger(tmpdir, version=logger.version) trainer_options['logger'] = new_logger trainer_options['checkpoint_callback'] = ModelCheckpoint(tmpdir) - trainer_options['train_percent_check'] = 0.2 + trainer_options['train_percent_check'] = 0.5 trainer_options['val_percent_check'] = 0.2 trainer_options['max_epochs'] = 1 new_trainer = Trainer(**trainer_options) @@ -269,68 +274,6 @@ def test_dp_resume(tmpdir): model.unfreeze() -def test_cpu_restore_training(tmpdir): - """Verify continue training session on CPU.""" - tutils.reset_seed() - - hparams = tutils.get_hparams() - model = LightningTestModel(hparams) - - # logger file to get meta - test_logger_version = 10 - logger = tutils.get_test_tube_logger(tmpdir, False, version=test_logger_version) - - trainer_options = dict( - max_epochs=8, - val_check_interval=0.50, - val_percent_check=0.2, - train_percent_check=0.2, - logger=logger, - checkpoint_callback=ModelCheckpoint(tmpdir, save_top_k=-1) - ) - - # fit model - trainer = Trainer(**trainer_options) - result = trainer.fit(model) - # Increment since we've finished the current epoch, don't want to rerun - real_global_epoch = trainer.current_epoch + 1 - - # traning complete - assert result == 1, 'amp + ddp model failed to complete' - - # 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 - new_logger = tutils.get_test_tube_logger(tmpdir, False, version=test_logger_version) - trainer_options = dict( - max_epochs=2, - val_check_interval=0.50, - val_percent_check=0.2, - train_percent_check=0.2, - logger=new_logger, - checkpoint_callback=ModelCheckpoint(tmpdir), - ) - 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_good_acc(): - assert trainer.current_epoch == real_global_epoch - assert trainer.current_epoch >= 0 - - # if model and state loaded correctly, predictions will be good even though we - # haven't trained with the new loaded model - trainer.model.eval() - for dataloader in trainer.val_dataloaders: - tutils.run_prediction(dataloader, trainer.model) - - model.on_train_start = assert_good_acc - - # 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) - - def test_model_saving_loading(tmpdir): """Tests use case where trainer saves the model, and user loads it from tags independently.""" tutils.reset_seed()