From 09167efdb59e1be8ffe9ff7010393bff084390be Mon Sep 17 00:00:00 2001 From: Jirka Borovec Date: Tue, 31 Mar 2020 00:37:02 +0200 Subject: [PATCH] Checkpointing interval (#1272) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * formatting * formatting * fix interval * fix train loop * fix test * parametrize test * Apply suggestions from code review Co-Authored-By: Adrian Wälchli * fix calling * flake8 * add types Co-authored-by: Adrian Wälchli Co-authored-by: William Falcon --- .../multi_node_ddp_demo.py | 6 +- pytorch_lightning/callbacks/base.py | 2 +- .../callbacks/model_checkpoint.py | 82 +++++----- pytorch_lightning/profiler/profiler.py | 39 ++--- .../trainer/distrib_data_parallel.py | 7 +- pytorch_lightning/trainer/distrib_parts.py | 24 ++- pytorch_lightning/trainer/evaluation_loop.py | 6 +- pytorch_lightning/trainer/trainer.py | 22 ++- pytorch_lightning/trainer/training_io.py | 25 +-- pytorch_lightning/trainer/training_loop.py | 33 ++-- tests/base/__init__.py | 4 +- tests/base/mixins.py | 11 +- tests/base/models.py | 29 +--- tests/test_profiler.py | 24 ++- tests/trainer/test_trainer.py | 142 +++--------------- 15 files changed, 162 insertions(+), 294 deletions(-) diff --git a/pl_examples/multi_node_examples/multi_node_ddp_demo.py b/pl_examples/multi_node_examples/multi_node_ddp_demo.py index fb28c060..eb04611b 100644 --- a/pl_examples/multi_node_examples/multi_node_ddp_demo.py +++ b/pl_examples/multi_node_examples/multi_node_ddp_demo.py @@ -16,11 +16,7 @@ np.random.seed(SEED) def main(hparams): - """ - Main training routine specific for this project - :param hparams: - :return: - """ + """Main training routine specific for this project.""" # ------------------------ # 1 INIT LIGHTNING MODEL # ------------------------ diff --git a/pytorch_lightning/callbacks/base.py b/pytorch_lightning/callbacks/base.py index 003241b2..c7898e9d 100644 --- a/pytorch_lightning/callbacks/base.py +++ b/pytorch_lightning/callbacks/base.py @@ -1,6 +1,6 @@ r""" Callback Base -============== +============= Abstract base class used to build new callbacks. """ diff --git a/pytorch_lightning/callbacks/model_checkpoint.py b/pytorch_lightning/callbacks/model_checkpoint.py index 981e07b3..54997e9e 100644 --- a/pytorch_lightning/callbacks/model_checkpoint.py +++ b/pytorch_lightning/callbacks/model_checkpoint.py @@ -94,7 +94,7 @@ class ModelCheckpoint(Callback): self.save_top_k = save_top_k self.save_weights_only = save_weights_only self.period = period - self.epochs_since_last_check = 0 + self.epoch_last_check = None self.prefix = prefix self.best_k_models = {} # {filename: monitor} @@ -139,21 +139,20 @@ class ModelCheckpoint(Callback): def format_checkpoint_name(self, epoch, metrics, ver=None): """Generate a filename according define template. - Examples - -------- - >>> tmpdir = os.path.dirname(__file__) - >>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{epoch}')) - >>> os.path.basename(ckpt.format_checkpoint_name(0, {})) - 'epoch=0.ckpt' - >>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{epoch:03d}')) - >>> os.path.basename(ckpt.format_checkpoint_name(5, {})) - 'epoch=005.ckpt' - >>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{epoch}-{val_loss:.2f}')) - >>> os.path.basename(ckpt.format_checkpoint_name(2, dict(val_loss=0.123456))) - 'epoch=2-val_loss=0.12.ckpt' - >>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{missing:d}')) - >>> os.path.basename(ckpt.format_checkpoint_name(0, {})) - 'missing=0.ckpt' + Examples: + >>> tmpdir = os.path.dirname(__file__) + >>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{epoch}')) + >>> os.path.basename(ckpt.format_checkpoint_name(0, {})) + 'epoch=0.ckpt' + >>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{epoch:03d}')) + >>> os.path.basename(ckpt.format_checkpoint_name(5, {})) + 'epoch=005.ckpt' + >>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{epoch}-{val_loss:.2f}')) + >>> os.path.basename(ckpt.format_checkpoint_name(2, dict(val_loss=0.123456))) + 'epoch=2-val_loss=0.12.ckpt' + >>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{missing:d}')) + >>> os.path.basename(ckpt.format_checkpoint_name(0, {})) + 'missing=0.ckpt' """ # check if user passed in keys to the string groups = re.findall(r'(\{.*?)[:\}]', self.filename) @@ -181,41 +180,36 @@ class ModelCheckpoint(Callback): metrics = trainer.callback_metrics epoch = trainer.current_epoch - self.epochs_since_last_check += 1 - if self.save_top_k == 0: # no models are saved return - if self.epochs_since_last_check >= self.period: - self.epochs_since_last_check = 0 + if self.epoch_last_check is not None and (epoch - self.epoch_last_check) < self.period: + # skipping in this term + return - filepath = self.format_checkpoint_name(epoch, metrics) - version_cnt = 0 - while os.path.isfile(filepath): - filepath = self.format_checkpoint_name(epoch, metrics, ver=version_cnt) - # this epoch called before - version_cnt += 1 + self.epoch_last_check = epoch - if self.save_top_k != -1: - current = metrics.get(self.monitor) + filepath = self.format_checkpoint_name(epoch, metrics) + version_cnt = 0 + while os.path.isfile(filepath): + filepath = self.format_checkpoint_name(epoch, metrics, ver=version_cnt) + # this epoch called before + version_cnt += 1 - if current is None: - warnings.warn( - f'Can save best model only with {self.monitor} available,' - ' skipping.', RuntimeWarning) - else: - if self.check_monitor_top_k(current): - self._do_check_save(filepath, current, epoch) - else: - if self.verbose > 0: - log.info( - f'\nEpoch {epoch:05d}: {self.monitor}' - f' was not in top {self.save_top_k}') + if self.save_top_k != -1: + current = metrics.get(self.monitor) - else: - if self.verbose > 0: - log.info(f'\nEpoch {epoch:05d}: saving model to {filepath}') - self._save_model(filepath) + if current is None: + warnings.warn(f'Can save best model only with {self.monitor} available, skipping.', RuntimeWarning) + elif self.check_monitor_top_k(current): + self._do_check_save(filepath, current, epoch) + elif self.verbose > 0: + log.info(f'\nEpoch {epoch:05d}: {self.monitor} was not in top {self.save_top_k}') + + else: + if self.verbose > 0: + log.info(f'\nEpoch {epoch:05d}: saving model to {filepath}') + self._save_model(filepath) def _do_check_save(self, filepath, current, epoch): # remove kth diff --git a/pytorch_lightning/profiler/profiler.py b/pytorch_lightning/profiler/profiler.py index 0b04f750..e2a6e5b2 100644 --- a/pytorch_lightning/profiler/profiler.py +++ b/pytorch_lightning/profiler/profiler.py @@ -17,15 +17,15 @@ class BaseProfiler(ABC): """ @abstractmethod - def start(self, action_name): + def start(self, action_name: str) -> None: """Defines how to start recording an action.""" @abstractmethod - def stop(self, action_name): + def stop(self, action_name: str) -> None: """Defines how to record the duration once an action is complete.""" @contextmanager - def profile(self, action_name): + def profile(self, action_name: str) -> None: """ Yields a context manager to encapsulate the scope of a profiled action. @@ -43,7 +43,7 @@ class BaseProfiler(ABC): finally: self.stop(action_name) - def profile_iterable(self, iterable, action_name): + def profile_iterable(self, iterable, action_name: str) -> None: iterator = iter(iterable) while True: try: @@ -55,7 +55,7 @@ class BaseProfiler(ABC): self.stop(action_name) break - def describe(self): + def describe(self) -> None: """Logs a profile report after the conclusion of the training run.""" pass @@ -69,10 +69,10 @@ class PassThroughProfiler(BaseProfiler): def __init__(self): pass - def start(self, action_name): + def start(self, action_name: str) -> None: pass - def stop(self, action_name): + def stop(self, action_name: str) -> None: pass @@ -86,14 +86,14 @@ class Profiler(BaseProfiler): self.current_actions = {} self.recorded_durations = defaultdict(list) - def start(self, action_name): + def start(self, action_name: str) -> None: if action_name in self.current_actions: raise ValueError( f"Attempted to start {action_name} which has already started." ) self.current_actions[action_name] = time.monotonic() - def stop(self, action_name): + def stop(self, action_name: str) -> None: end_time = time.monotonic() if action_name not in self.current_actions: raise ValueError( @@ -103,7 +103,7 @@ class Profiler(BaseProfiler): duration = end_time - start_time self.recorded_durations[action_name].append(duration) - def describe(self): + def describe(self) -> None: output_string = "\n\nProfiler Report\n" def log_row(action, mean, total): @@ -126,24 +126,25 @@ class AdvancedProfiler(BaseProfiler): verbose and you should only use this if you want very detailed reports. """ - def __init__(self, output_filename=None, line_count_restriction=1.0): + def __init__(self, output_filename: str = None, line_count_restriction: float = 1.0): """ - :param output_filename (str): optionally save profile results to file instead of printing - to std out when training is finished. - :param line_count_restriction (int|float): this can be used to limit the number of functions - reported for each action. either an integer (to select a count of lines), - or a decimal fraction between 0.0 and 1.0 inclusive (to select a percentage of lines) + Args: + output_filename: optionally save profile results to file instead of printing + to std out when training is finished. + line_count_restriction: this can be used to limit the number of functions + reported for each action. either an integer (to select a count of lines), + or a decimal fraction between 0.0 and 1.0 inclusive (to select a percentage of lines) """ self.profiled_actions = {} self.output_filename = output_filename self.line_count_restriction = line_count_restriction - def start(self, action_name): + def start(self, action_name: str) -> None: if action_name not in self.profiled_actions: self.profiled_actions[action_name] = cProfile.Profile() self.profiled_actions[action_name].enable() - def stop(self, action_name): + def stop(self, action_name: str) -> None: pr = self.profiled_actions.get(action_name) if pr is None: raise ValueError( # pragma: no-cover @@ -151,7 +152,7 @@ class AdvancedProfiler(BaseProfiler): ) pr.disable() - def describe(self): + def describe(self) -> None: self.recorded_stats = {} for action_name, pr in self.profiled_actions.items(): s = io.StringIO() diff --git a/pytorch_lightning/trainer/distrib_data_parallel.py b/pytorch_lightning/trainer/distrib_data_parallel.py index b90d089f..182b5404 100644 --- a/pytorch_lightning/trainer/distrib_data_parallel.py +++ b/pytorch_lightning/trainer/distrib_data_parallel.py @@ -199,10 +199,9 @@ class TrainerDDPMixin(ABC): self.use_ddp2 = distributed_backend == 'ddp2' elif distributed_backend is None: - m = 'You requested multiple GPUs but did not specify a backend' \ - 'Trainer(distributed_backend=dp) (or ddp, ddp2)' \ - 'Setting distributed_backend=dp for you' - warnings.warn(m) + warnings.warn('You requested multiple GPUs but did not specify a backend, e.g.' + ' Trainer(distributed_backend=dp) (or ddp, ddp2).' + ' Setting distributed_backend=dp for you.') self.use_dp = True self.use_ddp = False self.use_ddp2 = False diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py index 3eee3330..4718407a 100644 --- a/pytorch_lightning/trainer/distrib_parts.py +++ b/pytorch_lightning/trainer/distrib_parts.py @@ -491,9 +491,8 @@ class TrainerDPMixin(ABC): if self.precision == 16: os.environ['XLA_USE_BF16'] = str(1) - m = f'INIT TPU local core: {self.tpu_local_core_rank}, ' \ - f'global rank: {self.tpu_global_core_rank}' - log.info(m) + log.info(f'INIT TPU local core: {self.tpu_local_core_rank},' + f' global rank: {self.tpu_global_core_rank}') # continue training routine self.run_pretrain_routine(model) @@ -512,12 +511,10 @@ class TrainerDPMixin(ABC): # https://github.com/NVIDIA/apex/issues/227 if self.use_dp and self.use_amp: if self.amp_level == 'O2': - m = f""" - Amp level {self.amp_level} with DataParallel is not supported. - See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227. - We recommend you switch to ddp if you want to use amp - """ - raise MisconfigurationException(m) + raise MisconfigurationException( + f'Amp level {self.amp_level} with DataParallel is not supported.' + f' See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227.' + f' We recommend you switch to ddp if you want to use amp') else: model, optimizers = model.configure_apex(amp, model, self.optimizers, self.amp_level) @@ -584,11 +581,10 @@ def sanitize_gpu_ids(gpus): all_available_gpus = get_all_available_gpus() for gpu in gpus: if gpu not in all_available_gpus: - message = f""" - You requested GPUs: {gpus} - But your machine only has: {all_available_gpus} - """ - raise MisconfigurationException(message) + raise MisconfigurationException(f""" + You requested GPUs: {gpus} + But your machine only has: {all_available_gpus} + """) return gpus diff --git a/pytorch_lightning/trainer/evaluation_loop.py b/pytorch_lightning/trainer/evaluation_loop.py index 4e9aed6d..d720f81e 100644 --- a/pytorch_lightning/trainer/evaluation_loop.py +++ b/pytorch_lightning/trainer/evaluation_loop.py @@ -322,9 +322,9 @@ class TrainerEvaluationLoopMixin(ABC): def run_evaluation(self, test_mode: bool = False): # when testing make sure user defined a test step if test_mode and not self.is_overriden('test_step'): - m = "You called `.test()` without defining model's `.test_step()`." \ - " Please define and try again" - raise MisconfigurationException(m) + raise MisconfigurationException( + "You called `.test()` without defining model's `.test_step()`." + " Please define and try again") # Validation/Test begin callbacks if test_mode: diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index 6ba50bf2..6b4b461a 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -328,11 +328,8 @@ class Trainer( if self.fast_dev_run: self.num_sanity_val_steps = 1 self.max_epochs = 1 - m = ''' - Running in fast_dev_run mode: will run a full train, - val loop using a single batch - ''' - log.info(m) + log.info('Running in fast_dev_run mode: will run a full train,' + ' val loop using a single batch') # set default save path if user didn't provide one self.default_save_path = default_save_path @@ -739,22 +736,22 @@ class Trainer( # functions to overwrite with these implementations if train_dataloader is not None: if not self.is_overriden('training_step', model): - m = 'You called .fit() with a train_dataloader but did not define training_step()' - raise MisconfigurationException(m) + raise MisconfigurationException( + 'You called `.fit()` with a `train_dataloader` but did not define `training_step()`') model.train_dataloader = _PatchDataLoader(train_dataloader) if val_dataloaders is not None: if not self.is_overriden('validation_step', model): - m = 'You called .fit() with a val_dataloaders but did not define validation_step()' - raise MisconfigurationException(m) + raise MisconfigurationException( + 'You called `.fit()` with a `val_dataloaders` but did not define `validation_step()`') model.val_dataloader = _PatchDataLoader(val_dataloaders) if test_dataloaders is not None: if not self.is_overriden('test_step', model): - m = 'You called .fit() with a test_dataloaders but did not define test_step()' - raise MisconfigurationException(m) + raise MisconfigurationException( + 'You called `.fit()` with a `test_dataloaders` but did not define `test_step()`') model.test_dataloader = _PatchDataLoader(test_dataloaders) @@ -855,8 +852,7 @@ class Trainer( if self.weights_summary in ['full', 'top']: ref_model.summarize(mode=self.weights_summary) else: - m = "weights_summary can be None, 'full' or 'top'" - raise MisconfigurationException(m) + raise MisconfigurationException("weights_summary can be None, 'full' or 'top'") # track model now. # if cluster resets state, the model will update with the saved weights diff --git a/pytorch_lightning/trainer/training_io.py b/pytorch_lightning/trainer/training_io.py index 40238c27..636491f6 100644 --- a/pytorch_lightning/trainer/training_io.py +++ b/pytorch_lightning/trainer/training_io.py @@ -143,15 +143,12 @@ class TrainerIOMixin(ABC): # -------------------- # CHECK-POINTING # -------------------- - def restore_weights(self, model): + def restore_weights(self, model: LightningModule): """ 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 - - :param model: - :return: """ # clear cache before restore if self.on_gpu: @@ -230,17 +227,17 @@ class TrainerIOMixin(ABC): # -------------------- # MODEL SAVE CHECKPOINT # -------------------- - def _atomic_save(self, checkpoint, filepath): + def _atomic_save(self, checkpoint, filepath: str): """Saves a checkpoint atomically, avoiding the creation of incomplete checkpoints. This will create a temporary checkpoint with a suffix of ``.part``, then copy it to the final location once saving is finished. Args: - checkpoint (object): The object to save. + checkpoint: The object to save. Built to be used with the ``dump_checkpoint`` method, but can deal with anything which ``torch.save`` accepts. - filepath (str|pathlib.Path): The path to which the checkpoint will be saved. + filepath: The path to which the checkpoint will be saved. This points to the file that the checkpoint will be stored in. """ tmp_path = str(filepath) + ".part" @@ -260,7 +257,7 @@ class TrainerIOMixin(ABC): self._atomic_save(checkpoint, filepath) - def restore(self, checkpoint_path, on_gpu): + def restore(self, checkpoint_path: str, on_gpu: bool): """ Restore training state from checkpoint. Also restores all training state like: @@ -268,10 +265,6 @@ class TrainerIOMixin(ABC): - callbacks - schedulers - optimizer - :param checkpoint_path: - :param on_gpu: - - :return: """ # if on_gpu: @@ -341,12 +334,8 @@ class TrainerIOMixin(ABC): # -------------------- # HPC IO # -------------------- - def restore_hpc_weights_if_needed(self, model): - """ - If there is a set of hpc weights, use as signal to restore model - :param model: - :return: - """ + def restore_hpc_weights_if_needed(self, model: LightningModule): + """If there is a set of hpc weights, use as signal to restore model.""" did_restore = False # look for hpc weights diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py index 71896f8f..d0a8ee1e 100644 --- a/pytorch_lightning/trainer/training_loop.py +++ b/pytorch_lightning/trainer/training_loop.py @@ -461,8 +461,7 @@ class TrainerTrainLoopMixin(ABC): # CHECKPOINTING, EARLY STOPPING # --------------- # save checkpoint even when no test or val step are defined - train_step_only = not self.is_overriden('validation_step') - if self.fast_dev_run or should_check_val or train_step_only: + if self.fast_dev_run or should_check_val: self.call_checkpoint_callback() if self.enable_early_stop: @@ -483,6 +482,13 @@ class TrainerTrainLoopMixin(ABC): if early_stop_epoch or self.fast_dev_run: break + # in case validation step is missing and you are not running fast-dev to duplicate last batch + if not self.is_overriden('validation_step') and not (self.fast_dev_run or should_check_val): + self.call_checkpoint_callback() + + if self.enable_early_stop: + self.early_stop_callback.check_metrics(self.callback_metrics) + # Epoch end events with self.profiler.profile('on_epoch_end'): # callbacks @@ -709,20 +715,20 @@ class TrainerTrainLoopMixin(ABC): with self.profiler.profile('training_end'): output = model_ref.training_end(output) - m = 'training_end was deprecated in 0.7.0 and will be removed 1.0.0. ' \ - 'Use training_epoch_end instead' - warnings.warn(m, DeprecationWarning) + warnings.warn('`training_end` was deprecated in 0.7.0 and will be removed 1.0.0.' + ' Use training_epoch_end instead', DeprecationWarning) # format and reduce outputs accordingly output = self.process_output(output, train=True) return output - def update_learning_rates(self, interval): - ''' Update learning rates + def update_learning_rates(self, interval: str): + """Update learning rates. + Args: - interval (str): either 'epoch' or 'step'. - ''' + interval: either 'epoch' or 'step'. + """ if not self.lr_schedulers: return @@ -738,10 +744,11 @@ class TrainerTrainLoopMixin(ABC): monitor_val = self.callback_metrics.get(monitor_key) if monitor_val is None: avail_metrics = ','.join(list(self.callback_metrics.keys())) - m = f'ReduceLROnPlateau conditioned on metric {monitor_key} ' \ - f'which is not available. Available metrics are: {avail_metrics}. ' \ - 'Condition can be set using `monitor` key in lr scheduler dict' - raise MisconfigurationException(m) + raise MisconfigurationException( + f'ReduceLROnPlateau conditioned on metric {monitor_key}' + f' which is not available. Available metrics are: {avail_metrics}.' + ' Condition can be set using `monitor` key in lr scheduler dict' + ) lr_scheduler['scheduler'].step(monitor_val) else: lr_scheduler['scheduler'].step() diff --git a/tests/base/__init__.py b/tests/base/__init__.py index 1e684698..74157880 100644 --- a/tests/base/__init__.py +++ b/tests/base/__init__.py @@ -40,7 +40,7 @@ class LightningTestModel(LightTrainDataloader, class LightningTestModelWithoutHyperparametersArg(LightningTestModel): - """ without hparams argument in constructor """ + """Without hparams argument in constructor """ def __init__(self): import tests.base.utils as tutils @@ -51,7 +51,7 @@ class LightningTestModelWithoutHyperparametersArg(LightningTestModel): class LightningTestModelWithUnusedHyperparametersArg(LightningTestModelWithoutHyperparametersArg): - """ has hparams argument in constructor but is not used """ + """It has hparams argument in constructor but is not used.""" def __init__(self, hparams): super().__init__() diff --git a/tests/base/mixins.py b/tests/base/mixins.py index 1a05049f..7b40184a 100644 --- a/tests/base/mixins.py +++ b/tests/base/mixins.py @@ -14,11 +14,7 @@ class LightValidationStepMixin: return self._dataloader(train=False) def validation_step(self, batch, batch_idx, *args, **kwargs): - """ - Lightning calls this inside the validation loop - :param batch: - :return: - """ + """Lightning calls this inside the validation loop.""" x, y = batch x = x.view(x.size(0), -1) y_hat = self(x) @@ -66,8 +62,9 @@ class LightValidationMixin(LightValidationStepMixin): def validation_epoch_end(self, outputs): """ Called at the end of validation to aggregate outputs - :param outputs: list of individual outputs of each validation step - :return: + + Args: + outputs: list of individual outputs of each validation step """ # if returned a scalar from validation_step, outputs is a list of tensor scalars # we return just the average in this case (if we want) diff --git a/tests/base/models.py b/tests/base/models.py index 00baed58..b796b60c 100644 --- a/tests/base/models.py +++ b/tests/base/models.py @@ -42,16 +42,10 @@ class DictHparamsModel(LightningModule): class TestModelBase(LightningModule): - """ - Base LightningModule for testing. Implements only the required - interface - """ + """Base LightningModule for testing. Implements only the required interface.""" - def __init__(self, hparams, force_remove_distributed_sampler=False): - """ - Pass in parsed HyperOptArgumentParser to the model - :param hparams: - """ + def __init__(self, hparams, force_remove_distributed_sampler: bool = False): + """Pass in parsed HyperOptArgumentParser to the model.""" # init superclass super().__init__() self.hparams = hparams @@ -71,10 +65,7 @@ class TestModelBase(LightningModule): # MODEL SETUP # --------------------- def __build_model(self): - """ - Layout model - :return: - """ + """Layout model.""" 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) @@ -87,11 +78,7 @@ class TestModelBase(LightningModule): # TRAINING # --------------------- def forward(self, x): - """ - No special modification required for lightning, define as you normally would - :param x: - :return: - """ + """No special modification required for lightning, define as you normally would.""" x = self.c_d1(x) x = torch.tanh(x) x = self.c_d1_bn(x) @@ -107,11 +94,7 @@ class TestModelBase(LightningModule): return nll def training_step(self, batch, batch_idx, optimizer_idx=None): - """ - Lightning calls this inside the training loop - :param batch: - :return: - """ + """Lightning calls this inside the training loop""" # forward pass x, y = batch x = x.view(x.size(0), -1) diff --git a/tests/test_profiler.py b/tests/test_profiler.py index eae80bd6..e60476bc 100644 --- a/tests/test_profiler.py +++ b/tests/test_profiler.py @@ -35,7 +35,11 @@ def advanced_profiler(): return profiler -@pytest.mark.parametrize("action,expected", [("a", [3, 1]), ("b", [2]), ("c", [1])]) +@pytest.mark.parametrize(["action", "expected"], [ + pytest.param("a", [3, 1]), + pytest.param("b", [2]), + pytest.param("c", [1]) +]) def test_simple_profiler_durations(simple_profiler, action, expected): """Ensure the reported durations are reasonably accurate.""" @@ -50,7 +54,11 @@ def test_simple_profiler_durations(simple_profiler, action, expected): ) -@pytest.mark.parametrize("action,expected", [("a", [3, 1]), ("b", [2]), ("c", [1])]) +@pytest.mark.parametrize(["action", "expected"], [ + pytest.param("a", [3, 1]), + pytest.param("b", [2]), + pytest.param("c", [1]) +]) def test_simple_profiler_iterable_durations(simple_profiler, action, expected): """Ensure the reported durations are reasonably accurate.""" iterable = _sleep_generator(expected) @@ -96,7 +104,11 @@ def test_simple_profiler_value_errors(simple_profiler): simple_profiler.stop(action) -@pytest.mark.parametrize("action,expected", [("a", [3, 1]), ("b", [2]), ("c", [1])]) +@pytest.mark.parametrize(["action", "expected"], [ + pytest.param("a", [3, 1]), + pytest.param("b", [2]), + pytest.param("c", [1]) +]) def test_advanced_profiler_durations(advanced_profiler, action, expected): for duration in expected: @@ -114,7 +126,11 @@ def test_advanced_profiler_durations(advanced_profiler, action, expected): ) -@pytest.mark.parametrize("action,expected", [("a", [3, 1]), ("b", [2]), ("c", [1])]) +@pytest.mark.parametrize(["action", "expected"], [ + pytest.param("a", [3, 1]), + pytest.param("b", [2]), + pytest.param("c", [1]) +]) def test_advanced_profiler_iterable_durations(advanced_profiler, action, expected): """Ensure the reported durations are reasonably accurate.""" iterable = _sleep_generator(expected) diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py index 89d849d7..295ea3bd 100644 --- a/tests/trainer/test_trainer.py +++ b/tests/trainer/test_trainer.py @@ -248,7 +248,19 @@ def test_dp_output_reduce(): assert reduced['b']['c'] == out['b']['c'] -def test_model_checkpoint_options(tmpdir): +@pytest.mark.parametrize(["save_top_k", "file_prefix", "expected_files"], [ + pytest.param(-1, '', {'epoch=4.ckpt', 'epoch=3.ckpt', 'epoch=2.ckpt', 'epoch=1.ckpt', 'epoch=0.ckpt'}, + id="CASE K=-1 (all)"), + pytest.param(1, 'test_prefix_', {'test_prefix_epoch=4.ckpt'}, + id="CASE K=1 (2.5, epoch 4)"), + pytest.param(2, '', {'epoch=4.ckpt', 'epoch=2.ckpt'}, + id="CASE K=2 (2.5 epoch 4, 2.8 epoch 2)"), + pytest.param(4, '', {'epoch=1.ckpt', 'epoch=4.ckpt', 'epoch=3.ckpt', 'epoch=2.ckpt'}, + id="CASE K=4 (save all 4 base)"), + pytest.param(3, '', {'epoch=2.ckpt', 'epoch=3.ckpt', 'epoch=4.ckpt'}, + id="CASE K=3 (save the 2nd, 3rd, 4th model)"), +]) +def test_model_checkpoint_options(tmpdir, save_top_k, file_prefix, expected_files): """Test ModelCheckpoint options.""" def mock_save_function(filepath): @@ -258,13 +270,9 @@ def test_model_checkpoint_options(tmpdir): _ = LightningTestModel(hparams) # simulated losses - save_dir = os.path.join(tmpdir, '1') - os.mkdir(save_dir) losses = [10, 9, 2.8, 5, 2.5] - # ----------------- - # CASE K=-1 (all) - checkpoint_callback = ModelCheckpoint(save_dir, save_top_k=-1, verbose=1) + checkpoint_callback = ModelCheckpoint(tmpdir, save_top_k=save_top_k, prefix=file_prefix, verbose=1) checkpoint_callback.save_function = mock_save_function trainer = Trainer() @@ -274,127 +282,13 @@ def test_model_checkpoint_options(tmpdir): trainer.callback_metrics = {'val_loss': loss} checkpoint_callback.on_validation_end(trainer, trainer.get_model()) - file_lists = set(os.listdir(save_dir)) + file_lists = set(os.listdir(tmpdir)) - assert len(file_lists) == len(losses), "Should save all models when save_top_k=-1" + assert len(file_lists) == len(expected_files), \ + "Should save %i models when save_top_k=%i" % (len(expected_files), save_top_k) # verify correct naming - for fname in {'epoch=4.ckpt', - 'epoch=3.ckpt', - 'epoch=2.ckpt', - 'epoch=1.ckpt', - 'epoch=0.ckpt'}: - assert fname in file_lists - - save_dir = os.path.join(tmpdir, '2') - os.mkdir(save_dir) - - # ----------------- - # CASE K=0 (none) - checkpoint_callback = ModelCheckpoint(save_dir, save_top_k=0, verbose=1) - checkpoint_callback.save_function = mock_save_function - trainer = Trainer() - - # emulate callback's calls during the training - for i, loss in enumerate(losses): - trainer.current_epoch = i - trainer.callback_metrics = {'val_loss': loss} - checkpoint_callback.on_validation_end(trainer, trainer.get_model()) - - file_lists = os.listdir(save_dir) - - assert len(file_lists) == 0, "Should save 0 models when save_top_k=0" - - save_dir = os.path.join(tmpdir, '3') - os.mkdir(save_dir) - - # ----------------- - # CASE K=1 (2.5, epoch 4) - checkpoint_callback = ModelCheckpoint(save_dir, save_top_k=1, verbose=1, prefix='test_prefix_') - checkpoint_callback.save_function = mock_save_function - trainer = Trainer() - - # emulate callback's calls during the training - for i, loss in enumerate(losses): - trainer.current_epoch = i - trainer.callback_metrics = {'val_loss': loss} - checkpoint_callback.on_validation_end(trainer, trainer.get_model()) - - file_lists = set(os.listdir(save_dir)) - - assert len(file_lists) == 1, "Should save 1 model when save_top_k=1" - assert 'test_prefix_epoch=4.ckpt' in file_lists - - save_dir = os.path.join(tmpdir, '4') - os.mkdir(save_dir) - - # ----------------- - # CASE K=2 (2.5 epoch 4, 2.8 epoch 2) - # make sure other files don't get deleted - - checkpoint_callback = ModelCheckpoint(save_dir, save_top_k=2, verbose=1) - open(f"{save_dir}/other_file.ckpt", 'a').close() - checkpoint_callback.save_function = mock_save_function - trainer = Trainer() - - # emulate callback's calls during the training - for i, loss in enumerate(losses): - trainer.current_epoch = i - trainer.callback_metrics = {'val_loss': loss} - checkpoint_callback.on_validation_end(trainer, trainer.get_model()) - - file_lists = set(os.listdir(save_dir)) - - assert len(file_lists) == 3, 'Should save 2 model when save_top_k=2' - for fname in {'epoch=4.ckpt', - 'epoch=2.ckpt', - 'other_file.ckpt'}: - assert fname in file_lists - - save_dir = os.path.join(tmpdir, '5') - os.mkdir(save_dir) - - # ----------------- - # CASE K=4 (save all 4 base) - # multiple checkpoints within same epoch - - checkpoint_callback = ModelCheckpoint(save_dir, save_top_k=4, verbose=1) - checkpoint_callback.save_function = mock_save_function - trainer = Trainer() - - # emulate callback's calls during the training - for loss in losses: - trainer.current_epoch = 0 - trainer.callback_metrics = {'val_loss': loss} - checkpoint_callback.on_validation_end(trainer, trainer.get_model()) - - file_lists = set(os.listdir(save_dir)) - - assert len(file_lists) == 4, 'Should save all 4 models when save_top_k=4 within same epoch' - - save_dir = os.path.join(tmpdir, '6') - os.mkdir(save_dir) - - # ----------------- - # CASE K=3 (save the 2nd, 3rd, 4th model) - # multiple checkpoints within same epoch - - checkpoint_callback = ModelCheckpoint(save_dir, save_top_k=3, verbose=1) - checkpoint_callback.save_function = mock_save_function - trainer = Trainer() - - # emulate callback's calls during the training - for loss in losses: - trainer.current_epoch = 0 - trainer.callback_metrics = {'val_loss': loss} - checkpoint_callback.on_validation_end(trainer, trainer.get_model()) - - file_lists = set(os.listdir(save_dir)) - - assert len(file_lists) == 3, 'Should save 3 models when save_top_k=3' - for fname in {'epoch=0.ckpt', - 'epoch=0.ckpt', - 'epoch=0.ckpt'}: + for fname in expected_files: assert fname in file_lists